commit 59a0a3844cfa95119968e10e12a74b135438f1c4 Author: wehub-resource-sync Date: Mon Jul 13 12:32:31 2026 +0800 chore: import upstream snapshot with attribution diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000..a06cb67 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,11 @@ +# Skip pre-compiled kernel .so/.o files so they are rebuilt inside the container +**/*.so +**/*.o + +# Common development artifacts +**/__pycache__ +*.pyc +.git +*.nsys-rep +*.sqlite +*.egg-info diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS new file mode 100644 index 0000000..f994027 --- /dev/null +++ b/.github/CODEOWNERS @@ -0,0 +1,5 @@ +/python/tokenspeed/runtime @lightseekorg/code-owner +/tokenspeed-kernel @lightseekorg/code-owner @lightseekorg/kernel-team +/tokenspeed-kernel-amd @lightseekorg/code-owner @lightseekorg/kernel-team +/tokenspeed-mla @lightseekorg/code-owner +/tokenspeed-scheduler @lightseekorg/code-owner diff --git a/.github/ISSUE_TEMPLATE/1-bug-report.yml b/.github/ISSUE_TEMPLATE/1-bug-report.yml new file mode 100644 index 0000000..1aceb69 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/1-bug-report.yml @@ -0,0 +1,47 @@ +name: 🐞 Bug report +description: Report a reproducible issue +title: "[Bug] " +labels: ['Bug'] + +body: +- type: checkboxes + attributes: + label: Checklist + options: + - label: "I searched existing issues before filing this report." + required: true + - label: "I am running on a supported GPU — NVIDIA: H100, H200, B200, B300, GB200, GB300; or AMD: MI350X, MI355X. (Other hardware is not supported and issues will be closed.)" + required: true + - label: "NVIDIA users: I am using CUDA 13 and CUDA driver ≥ 580 (recommended). If not, we may not be able to reproduce the issue." +- type: textarea + attributes: + label: Summary + description: What is the problem? Describe what you observed and what you expected. + validations: + required: true +- type: textarea + attributes: + label: Reproduction + description: Provide the server launch command and a minimal client command or script to reproduce the issue. + placeholder: | + Server: + ```bash + + ``` + + Client (minimal reproduction): + ```bash + + ``` + validations: + required: true +- type: textarea + attributes: + label: Environment + description: Paste the output of `tokenspeed env`. + placeholder: | + ``` + tokenspeed env + ``` + 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 0000000..d9a7492 --- /dev/null +++ b/.github/ISSUE_TEMPLATE/2-feature-request.yml @@ -0,0 +1,33 @@ +name: 🚀 Feature request +description: Propose a concrete improvement to TokenSpeed +title: "[Feature] " + +body: +- type: checkboxes + attributes: + label: Checklist + options: + - label: "I searched existing issues and open PRs before filing this request." + required: true +- type: textarea + attributes: + label: Problem and motivation + description: > + What online serving pain point or production need does this address? + If possible, share your current deployment scale and use case. + Help us understand why this is necessary and what the concrete benefit would be. + validations: + required: true +- type: textarea + attributes: + label: Proposed solution + description: > + RFC, design doc, or a description of the proposed solution. + If you have a concrete design in mind, share it here. +- type: textarea + attributes: + label: Execution plan + description: > + How do you plan to implement this? Note that core features are developed and driven by the TokenSpeed core team. + If after evaluation this is not on or planned for the roadmap, we encourage fork. + For technical discussion, feel free to join us at http://slack.lightseek.org. diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md new file mode 100644 index 0000000..54a9ac8 --- /dev/null +++ b/.github/pull_request_template.md @@ -0,0 +1,7 @@ +## Summary + + + +## Test Plan + + diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml new file mode 100644 index 0000000..0ace0f8 --- /dev/null +++ b/.github/workflows/docs.yml @@ -0,0 +1,63 @@ +name: Docs + +on: + pull_request: + branches: [main] + paths: + - "docs/**" + - "README.md" + - ".github/workflows/docs.yml" + push: + branches: [main] + paths: + - "docs/**" + - "README.md" + - ".github/workflows/docs.yml" + workflow_dispatch: + +permissions: + contents: read + pages: write + id-token: write + +concurrency: + group: docs-${{ github.ref }} + cancel-in-progress: true + +jobs: + build: + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Setup Bun + uses: oven-sh/setup-bun@v2 + with: + bun-version: "1.3.0" + + - name: Install docs dependencies + working-directory: docs + run: bun install --frozen-lockfile + + - name: Build docs + working-directory: docs + run: bun run docs:build + + - name: Upload pages artifact + if: github.event_name != 'pull_request' + uses: actions/upload-pages-artifact@v3 + with: + path: docs/.vitepress/dist + + deploy: + if: github.event_name != 'pull_request' && github.ref == 'refs/heads/main' && github.repository == 'lightseekorg/tokenspeed' + needs: build + runs-on: ubuntu-latest + environment: + name: github-pages + url: ${{ steps.deployment.outputs.page_url }} + steps: + - name: Deploy to GitHub Pages + id: deployment + uses: actions/deploy-pages@v4 diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml new file mode 100644 index 0000000..d6450e0 --- /dev/null +++ b/.github/workflows/lint.yml @@ -0,0 +1,94 @@ +name: Lint + +on: [pull_request] + +concurrency: + group: lint-${{ github.ref }} + cancel-in-progress: true + +jobs: + lint: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v2 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: '3.12' + + - name: Install pre-commit hook + run: | + python -m pip install pre-commit ruff "nanobind>=2.13.0" tokenspeed-spdlog==1.15.1 + pre-commit install + + - name: Linting + run: pre-commit run --all-files + + - name: Lint runtime Python with ruff + run: | + python -m ruff check \ + --select F401,F841,F821,F823,F811,F541,UP035,UP006,UP007,UP045 \ + python/tokenspeed/runtime + + - name: Install scheduler C++ lint tools + run: | + sudo apt-get update + sudo apt-get install -y clang-format clang-tidy cmake ninja-build + + - name: Configure scheduler compile database + run: | + cmake -S tokenspeed-scheduler -B tokenspeed-scheduler/build \ + -G Ninja \ + -DTOKENSPEED_SCHEDULER_BUILD_TESTS=ON \ + -DTOKENSPEED_SCHEDULER_BUILD_PYTHON=ON \ + -DCMAKE_EXPORT_COMPILE_COMMANDS=ON + + - name: Lint scheduler Python and file starts + run: | + python -m ruff check tokenspeed-scheduler + python - <<'PY' + import subprocess + from pathlib import Path + + suffixes = {".py", ".c", ".cc", ".cpp", ".cxx", ".h", ".hh", ".hpp", ".hxx"} + files = subprocess.check_output(["git", "ls-files", "tokenspeed-scheduler"], text=True).splitlines() + blank_starts = [] + relative_imports = [] + for name in files: + path = Path(name) + if path.suffix in suffixes and path.read_bytes().startswith((b"\n", b"\r\n")): + blank_starts.append(name) + if path.suffix == ".py": + for line_no, line in enumerate(path.read_text().splitlines(), 1): + if line.startswith(("from .", "import .")): + relative_imports.append(f"{name}:{line_no}: {line}") + + errors = [] + if blank_starts: + errors.append("files must not start with a blank line:\n" + "\n".join(blank_starts)) + if relative_imports: + errors.append("scheduler Python must use absolute imports:\n" + "\n".join(relative_imports)) + if errors: + raise SystemExit("\n\n".join(errors)) + PY + + - name: Lint scheduler C++ + run: | + cpp_files="$(git ls-files tokenspeed-scheduler | grep -E '\.(c|cc|cpp|cxx|h|hh|hpp|hxx)$')" + jobs="$(nproc)" + tidy_files="$(printf '%s\n' "$cpp_files" | grep -E '\.(c|cc|cpp|cxx)$' || true)" + format_count="$(printf '%s\n' "$cpp_files" | sed '/^$/d' | wc -l)" + tidy_count="$(printf '%s\n' "$tidy_files" | sed '/^$/d' | wc -l)" + + echo "Running clang-format on ${format_count} files with ${jobs} workers" + printf '%s\n' "$cpp_files" | xargs -r -n 16 -P "$jobs" sh -c ' + echo "clang-format: $*" + clang-format --dry-run --Werror "$@" + ' sh + + echo "Running clang-tidy on ${tidy_count} translation units with ${jobs} workers" + printf '%s\n' "$tidy_files" | xargs -r -n 1 -P "$jobs" sh -c ' + echo "clang-tidy: $1" + clang-tidy -p tokenspeed-scheduler/build "$1" + ' sh diff --git a/.github/workflows/pr-test-amd.yml b/.github/workflows/pr-test-amd.yml new file mode 100644 index 0000000..8f67005 --- /dev/null +++ b/.github/workflows/pr-test-amd.yml @@ -0,0 +1,241 @@ +name: PR Test AMD + +# Split vendor CI into separate workflows so either side can be rerun without +# waiting for the other side's jobs to finish. +on: + push: + branches: [main] + paths: + - "python/pyproject.toml" + - "python/tokenspeed/**" + - "tokenspeed-kernel/**" + - "tokenspeed-kernel-amd/**" + - "tokenspeed-mla/**" + - "tokenspeed-scheduler/**" + - "test/**" + - ".github/workflows/pr-test-amd.yml" + - ".github/workflows/pr-test-nvidia.yml" + - ".github/workflows/pr-test-nvidia-arm.yml" + pull_request: + types: [opened, synchronize, reopened, ready_for_review, closed] + branches: [main] + paths: + - "python/pyproject.toml" + - "python/tokenspeed/**" + - "tokenspeed-kernel/**" + - "tokenspeed-kernel-amd/**" + - "tokenspeed-mla/**" + - "tokenspeed-scheduler/**" + - "test/**" + - ".github/workflows/pr-test-amd.yml" + - ".github/workflows/pr-test-nvidia.yml" + - ".github/workflows/pr-test-nvidia-arm.yml" + workflow_dispatch: + inputs: + trigger: + description: "CI task trigger to run" + required: false + default: manual + type: choice + options: + - manual + - debug + +env: + PIP_BREAK_SYSTEM_PACKAGES: 1 + CUDA_VERSION: 13.0.1 + TOKENSPEED_B200_RUNNER_LABEL: ${{ vars.TOKENSPEED_B200_RUNNER_LABEL }} + +concurrency: + group: pr-test-amd-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: ${{ github.event_name != 'workflow_dispatch' }} + +jobs: + cancel-on-close: + if: github.event_name == 'pull_request' && github.event.action == 'closed' + runs-on: ubuntu-latest + steps: + - name: Cancel superseded PR Test runs + run: | + echo "Cancelling older PR Test AMD runs for PR #${{ github.event.pull_request.number }}" + + scan: + # Closed PRs only need cancel-on-close. For all other events, run jobs for + # pull_request events or for push/manual events in lightseekorg/tokenspeed. + if: | + (github.event_name != 'pull_request' || github.event.action != 'closed') && + ( + github.event_name == 'pull_request' || + github.repository == 'lightseekorg/tokenspeed' + ) + runs-on: ubuntu-latest + outputs: + matrix: ${{ steps.scan.outputs.matrix }} + install_tokenspeed_mla_from_source: ${{ steps.detect_mla.outputs.changed }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install scan dependency + # pypi.org occasionally returns "Content-Type: Unknown" for the simple + # index, which strict pip>=22 refuses to parse. Retry a few times. + run: | + for i in 1 2 3 4 5; do + python3 -m pip install pyyaml && exit 0 + echo "pyyaml install attempt $i failed, retrying in 10s..." + sleep 10 + done + exit 1 + + - name: Build task matrix + id: scan + run: | + MATRIX=$(python3 test/ci_system/pipeline.py scan \ + --root test/ci \ + --trigger "${{ github.event_name == 'workflow_dispatch' && github.event.inputs.trigger || 'per-commit' }}" \ + --runner-group "amd") + echo "matrix=${MATRIX}" >> "$GITHUB_OUTPUT" + + # Compute this once on the hosted runner (where `gh` is pre-installed) + # and propagate the result to every matrix job via job output. The + # downstream unit-test job reads it as an env var and `install_deps.sh` + # gates its in-tree `tokenspeed-mla` reinstall on it. Without this + # detection, CI would always exercise the pinned PyPI wheel and edits + # under `tokenspeed-mla/` would be silently uncovered. + - name: Detect tokenspeed-mla source changes + id: detect_mla + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + set -euo pipefail + case "${{ github.event_name }}" in + pull_request) + files=$(gh api \ + "/repos/${{ github.repository }}/pulls/${{ github.event.pull_request.number }}/files" \ + --paginate -q '.[].filename') + ;; + push) + files=$(gh api \ + "/repos/${{ github.repository }}/compare/${{ github.event.before }}...${{ github.event.after }}" \ + -q '.files[].filename') + ;; + *) + files="" + ;; + esac + if printf '%s\n' "$files" | grep -q '^tokenspeed-mla/'; then + echo "tokenspeed-mla touched by this run; in-tree source will be installed." + echo "changed=1" >> "$GITHUB_OUTPUT" + else + echo "tokenspeed-mla untouched; keeping pinned PyPI wheel." + echo "changed=0" >> "$GITHUB_OUTPUT" + fi + + unit-test: + needs: scan + if: | + (github.event_name != 'pull_request' || github.event.action != 'closed') && + ( + github.event_name == 'pull_request' || + github.repository == 'lightseekorg/tokenspeed' + ) && + ( + github.event_name != 'pull_request' || + github.event.pull_request.draft == false + ) + strategy: + fail-fast: false + matrix: ${{ fromJson(needs.scan.outputs.matrix) }} + name: ${{ matrix.name }} (${{ matrix.runner }}) + runs-on: ${{ matrix.runner }} + timeout-minutes: 60 + continue-on-error: ${{ matrix.optional }} + env: + # When the scan job determined that this run touched `tokenspeed-mla/`, + # `install_deps.sh` will reinstall the in-tree source on top of the + # pinned PyPI wheel. See the matching detect step in the `scan` job. + INSTALL_TOKENSPEED_MLA_FROM_SOURCE: ${{ needs.scan.outputs.install_tokenspeed_mla_from_source }} + # Authenticate Hugging Face Hub downloads so checkpoint pulls during + # `ts serve` startup are not throttled by the anonymous rate limit. The + # 600s per-file `runtime-1gpu` timeout is otherwise exceeded on cold + # runners (e.g. `amd-mi35x-1gpu-test` loading `openai/gpt-oss-120b` + + # EAGLE3 draft). + HF_TOKEN: ${{ secrets.HF_TOKEN }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + path: run-${{ github.run_id }}-${{ github.run_attempt }}-${{ matrix.name }}-${{ matrix.runner }} + + - name: Set work directory + run: | + echo "WORK_DIR=${{ github.workspace }}/run-${{ github.run_id }}-${{ github.run_attempt }}-${{ matrix.name }}-${{ matrix.runner }}" >> "$GITHUB_ENV" + + - name: Install pipeline dependency + # See note in the scan job — pypi.org "Content-Type: Unknown" flake. + run: | + for i in 1 2 3 4 5; do + python3 -m pip install pyyaml && exit 0 + echo "pyyaml install attempt $i failed, retrying in 10s..." + sleep 10 + done + exit 1 + + # When the diff touches `tokenspeed-mla/`, ask install_deps.sh to + # rebuild + install that package from the in-tree source after the + # pinned PyPI wheel has been installed (transitively via + # tokenspeed-kernel). Without this, CI keeps testing whichever + # `tokenspeed-mla` version is pinned in + # `tokenspeed-kernel/python/requirements/cuda-thirdparty.txt`, not + # the code on the PR/branch. When the diff doesn't touch + # `tokenspeed-mla/` we keep using the pinned PyPI wheel. + - name: Install CI task + run: | + cd "${{ env.WORK_DIR }}" + mkdir -p .ci-artifacts + python3 test/ci_system/pipeline.py execute \ + --config "${{ matrix.config }}" \ + --runner "${{ matrix.runner }}" \ + --work-dir "${{ env.WORK_DIR }}" \ + --print-plan \ + --only-stage install \ + --keep-runner-state \ + --result-json .ci-artifacts/result.json + + - name: Execute CI task + run: | + cd "${{ env.WORK_DIR }}" + mkdir -p .ci-artifacts + python3 test/ci_system/pipeline.py execute \ + --config "${{ matrix.config }}" \ + --runner "${{ matrix.runner }}" \ + --work-dir "${{ env.WORK_DIR }}" \ + --print-plan \ + --skip-stage install \ + --reuse-runner-state \ + --result-json .ci-artifacts/result.json + + - name: Upload task result + if: always() + uses: actions/upload-artifact@v4 + with: + name: pr-test-${{ matrix.name }}-${{ matrix.runner }}-${{ github.run_id }}-${{ github.run_attempt }} + path: ${{ env.WORK_DIR }}/.ci-artifacts/result.json + + - name: Cleanup work directory + if: always() + run: rm -rf "${{ env.WORK_DIR }}" + + finish: + if: | + always() && + (github.event_name != 'pull_request' || github.event.action != 'closed') && + ( + github.event_name == 'pull_request' || + github.repository == 'lightseekorg/tokenspeed' + ) + needs: [unit-test] + runs-on: ubuntu-latest + steps: + - name: Finish + run: echo "PR Test finished." diff --git a/.github/workflows/pr-test-nvidia-arm.yml b/.github/workflows/pr-test-nvidia-arm.yml new file mode 100644 index 0000000..494a097 --- /dev/null +++ b/.github/workflows/pr-test-nvidia-arm.yml @@ -0,0 +1,241 @@ +name: PR Test NVIDIA ARM + +# Split GB200/Grace NVIDIA CI into its own workflow so ARM jobs can be rerun +# without waiting for the other NVIDIA runner pools. +on: + push: + branches: [main] + paths: + - "python/pyproject.toml" + - "python/tokenspeed/**" + - "tokenspeed-kernel/**" + - "tokenspeed-kernel-amd/**" + - "tokenspeed-mla/**" + - "tokenspeed-scheduler/**" + - "test/**" + - ".github/workflows/pr-test-amd.yml" + - ".github/workflows/pr-test-nvidia.yml" + - ".github/workflows/pr-test-nvidia-arm.yml" + pull_request: + types: [opened, synchronize, reopened, ready_for_review, closed] + branches: [main] + paths: + - "python/pyproject.toml" + - "python/tokenspeed/**" + - "tokenspeed-kernel/**" + - "tokenspeed-kernel-amd/**" + - "tokenspeed-mla/**" + - "tokenspeed-scheduler/**" + - "test/**" + - ".github/workflows/pr-test-amd.yml" + - ".github/workflows/pr-test-nvidia.yml" + - ".github/workflows/pr-test-nvidia-arm.yml" + workflow_dispatch: + inputs: + trigger: + description: "CI task trigger to run" + required: false + default: manual + type: choice + options: + - manual + - debug + +env: + PIP_BREAK_SYSTEM_PACKAGES: 1 + CUDA_VERSION: 13.0.1 + TOKENSPEED_B200_RUNNER_LABEL: ${{ vars.TOKENSPEED_B200_RUNNER_LABEL }} + +concurrency: + group: pr-test-nvidia-arm-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: ${{ github.event_name != 'workflow_dispatch' }} + +jobs: + cancel-on-close: + if: github.event_name == 'pull_request' && github.event.action == 'closed' + runs-on: ubuntu-latest + steps: + - name: Cancel superseded PR Test runs + run: | + echo "Cancelling older PR Test NVIDIA ARM runs for PR #${{ github.event.pull_request.number }}" + + scan: + # Closed PRs only need cancel-on-close. For all other events, run jobs for + # pull_request events or for push/manual events in lightseekorg/tokenspeed. + if: | + (github.event_name != 'pull_request' || github.event.action != 'closed') && + ( + github.event_name == 'pull_request' || + github.repository == 'lightseekorg/tokenspeed' + ) + runs-on: ubuntu-latest + outputs: + matrix: ${{ steps.scan.outputs.matrix }} + install_tokenspeed_mla_from_source: ${{ steps.detect_mla.outputs.changed }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install scan dependency + # pypi.org occasionally returns "Content-Type: Unknown" for the simple + # index, which strict pip>=22 refuses to parse. Retry a few times. + run: | + for i in 1 2 3 4 5; do + python3 -m pip install pyyaml && exit 0 + echo "pyyaml install attempt $i failed, retrying in 10s..." + sleep 10 + done + exit 1 + + - name: Build task matrix + id: scan + run: | + MATRIX=$(python3 test/ci_system/pipeline.py scan \ + --root test/ci \ + --trigger "${{ github.event_name == 'workflow_dispatch' && github.event.inputs.trigger || 'per-commit' }}" \ + --runner-group "nvidia-arm") + echo "matrix=${MATRIX}" >> "$GITHUB_OUTPUT" + + # Compute this once on the hosted runner (where `gh` is pre-installed) + # and propagate the result to every matrix job via job output. The + # downstream unit-test job reads it as an env var and `install_deps.sh` + # gates its in-tree `tokenspeed-mla` reinstall on it. Without this + # detection, CI would always exercise the pinned PyPI wheel and edits + # under `tokenspeed-mla/` would be silently uncovered. + - name: Detect tokenspeed-mla source changes + id: detect_mla + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + set -euo pipefail + case "${{ github.event_name }}" in + pull_request) + files=$(gh api \ + "/repos/${{ github.repository }}/pulls/${{ github.event.pull_request.number }}/files" \ + --paginate -q '.[].filename') + ;; + push) + files=$(gh api \ + "/repos/${{ github.repository }}/compare/${{ github.event.before }}...${{ github.event.after }}" \ + -q '.files[].filename') + ;; + *) + files="" + ;; + esac + if printf '%s\n' "$files" | grep -q '^tokenspeed-mla/'; then + echo "tokenspeed-mla touched by this run; in-tree source will be installed." + echo "changed=1" >> "$GITHUB_OUTPUT" + else + echo "tokenspeed-mla untouched; keeping pinned PyPI wheel." + echo "changed=0" >> "$GITHUB_OUTPUT" + fi + + unit-test: + needs: scan + if: | + (github.event_name != 'pull_request' || github.event.action != 'closed') && + ( + github.event_name == 'pull_request' || + github.repository == 'lightseekorg/tokenspeed' + ) && + ( + github.event_name != 'pull_request' || + github.event.pull_request.draft == false + ) + strategy: + fail-fast: false + matrix: ${{ fromJson(needs.scan.outputs.matrix) }} + name: ${{ matrix.name }} (${{ matrix.runner }}) + runs-on: ${{ matrix.runner }} + timeout-minutes: 60 + continue-on-error: ${{ matrix.optional }} + env: + # When the scan job determined that this run touched `tokenspeed-mla/`, + # `install_deps.sh` will reinstall the in-tree source on top of the + # pinned PyPI wheel. See the matching detect step in the `scan` job. + INSTALL_TOKENSPEED_MLA_FROM_SOURCE: ${{ needs.scan.outputs.install_tokenspeed_mla_from_source }} + # Authenticate Hugging Face Hub downloads so checkpoint pulls during + # `ts serve` startup are not throttled by the anonymous rate limit. The + # 600s per-file `runtime-1gpu` timeout is otherwise exceeded on cold + # runners (e.g. `amd-mi35x-1gpu-test` loading `openai/gpt-oss-120b` + + # EAGLE3 draft). + HF_TOKEN: ${{ secrets.HF_TOKEN }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + path: run-${{ github.run_id }}-${{ github.run_attempt }}-${{ matrix.name }}-${{ matrix.runner }} + + - name: Set work directory + run: | + echo "WORK_DIR=${{ github.workspace }}/run-${{ github.run_id }}-${{ github.run_attempt }}-${{ matrix.name }}-${{ matrix.runner }}" >> "$GITHUB_ENV" + + - name: Install pipeline dependency + # See note in the scan job — pypi.org "Content-Type: Unknown" flake. + run: | + for i in 1 2 3 4 5; do + python3 -m pip install pyyaml && exit 0 + echo "pyyaml install attempt $i failed, retrying in 10s..." + sleep 10 + done + exit 1 + + # When the diff touches `tokenspeed-mla/`, ask install_deps.sh to + # rebuild + install that package from the in-tree source after the + # pinned PyPI wheel has been installed (transitively via + # tokenspeed-kernel). Without this, CI keeps testing whichever + # `tokenspeed-mla` version is pinned in + # `tokenspeed-kernel/python/requirements/cuda-thirdparty.txt`, not + # the code on the PR/branch. When the diff doesn't touch + # `tokenspeed-mla/` we keep using the pinned PyPI wheel. + - name: Install CI task + run: | + cd "${{ env.WORK_DIR }}" + mkdir -p .ci-artifacts + python3 test/ci_system/pipeline.py execute \ + --config "${{ matrix.config }}" \ + --runner "${{ matrix.runner }}" \ + --work-dir "${{ env.WORK_DIR }}" \ + --print-plan \ + --only-stage install \ + --keep-runner-state \ + --result-json .ci-artifacts/result.json + + - name: Execute CI task + run: | + cd "${{ env.WORK_DIR }}" + mkdir -p .ci-artifacts + python3 test/ci_system/pipeline.py execute \ + --config "${{ matrix.config }}" \ + --runner "${{ matrix.runner }}" \ + --work-dir "${{ env.WORK_DIR }}" \ + --print-plan \ + --skip-stage install \ + --reuse-runner-state \ + --result-json .ci-artifacts/result.json + + - name: Upload task result + if: always() + uses: actions/upload-artifact@v4 + with: + name: pr-test-${{ matrix.name }}-${{ matrix.runner }}-${{ github.run_id }}-${{ github.run_attempt }} + path: ${{ env.WORK_DIR }}/.ci-artifacts/result.json + + - name: Cleanup work directory + if: always() + run: rm -rf "${{ env.WORK_DIR }}" + + finish: + if: | + always() && + (github.event_name != 'pull_request' || github.event.action != 'closed') && + ( + github.event_name == 'pull_request' || + github.repository == 'lightseekorg/tokenspeed' + ) + needs: [unit-test] + runs-on: ubuntu-latest + steps: + - name: Finish + run: echo "PR Test NVIDIA ARM finished." diff --git a/.github/workflows/pr-test-nvidia.yml b/.github/workflows/pr-test-nvidia.yml new file mode 100644 index 0000000..452503d --- /dev/null +++ b/.github/workflows/pr-test-nvidia.yml @@ -0,0 +1,241 @@ +name: PR Test NVIDIA + +# Split vendor CI into separate workflows so either side can be rerun without +# waiting for the other side's jobs to finish. +on: + push: + branches: [main] + paths: + - "python/pyproject.toml" + - "python/tokenspeed/**" + - "tokenspeed-kernel/**" + - "tokenspeed-kernel-amd/**" + - "tokenspeed-mla/**" + - "tokenspeed-scheduler/**" + - "test/**" + - ".github/workflows/pr-test-amd.yml" + - ".github/workflows/pr-test-nvidia.yml" + - ".github/workflows/pr-test-nvidia-arm.yml" + pull_request: + types: [opened, synchronize, reopened, ready_for_review, closed] + branches: [main] + paths: + - "python/pyproject.toml" + - "python/tokenspeed/**" + - "tokenspeed-kernel/**" + - "tokenspeed-kernel-amd/**" + - "tokenspeed-mla/**" + - "tokenspeed-scheduler/**" + - "test/**" + - ".github/workflows/pr-test-amd.yml" + - ".github/workflows/pr-test-nvidia.yml" + - ".github/workflows/pr-test-nvidia-arm.yml" + workflow_dispatch: + inputs: + trigger: + description: "CI task trigger to run" + required: false + default: manual + type: choice + options: + - manual + - debug + +env: + PIP_BREAK_SYSTEM_PACKAGES: 1 + CUDA_VERSION: 13.0.1 + TOKENSPEED_B200_RUNNER_LABEL: ${{ vars.TOKENSPEED_B200_RUNNER_LABEL }} + +concurrency: + group: pr-test-nvidia-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: ${{ github.event_name != 'workflow_dispatch' }} + +jobs: + cancel-on-close: + if: github.event_name == 'pull_request' && github.event.action == 'closed' + runs-on: ubuntu-latest + steps: + - name: Cancel superseded PR Test runs + run: | + echo "Cancelling older PR Test NVIDIA runs for PR #${{ github.event.pull_request.number }}" + + scan: + # Closed PRs only need cancel-on-close. For all other events, run jobs for + # pull_request events or for push/manual events in lightseekorg/tokenspeed. + if: | + (github.event_name != 'pull_request' || github.event.action != 'closed') && + ( + github.event_name == 'pull_request' || + github.repository == 'lightseekorg/tokenspeed' + ) + runs-on: ubuntu-latest + outputs: + matrix: ${{ steps.scan.outputs.matrix }} + install_tokenspeed_mla_from_source: ${{ steps.detect_mla.outputs.changed }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Install scan dependency + # pypi.org occasionally returns "Content-Type: Unknown" for the simple + # index, which strict pip>=22 refuses to parse. Retry a few times. + run: | + for i in 1 2 3 4 5; do + python3 -m pip install pyyaml && exit 0 + echo "pyyaml install attempt $i failed, retrying in 10s..." + sleep 10 + done + exit 1 + + - name: Build task matrix + id: scan + run: | + MATRIX=$(python3 test/ci_system/pipeline.py scan \ + --root test/ci \ + --trigger "${{ github.event_name == 'workflow_dispatch' && github.event.inputs.trigger || 'per-commit' }}" \ + --runner-group "nvidia-x86") + echo "matrix=${MATRIX}" >> "$GITHUB_OUTPUT" + + # Compute this once on the hosted runner (where `gh` is pre-installed) + # and propagate the result to every matrix job via job output. The + # downstream unit-test job reads it as an env var and `install_deps.sh` + # gates its in-tree `tokenspeed-mla` reinstall on it. Without this + # detection, CI would always exercise the pinned PyPI wheel and edits + # under `tokenspeed-mla/` would be silently uncovered. + - name: Detect tokenspeed-mla source changes + id: detect_mla + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + set -euo pipefail + case "${{ github.event_name }}" in + pull_request) + files=$(gh api \ + "/repos/${{ github.repository }}/pulls/${{ github.event.pull_request.number }}/files" \ + --paginate -q '.[].filename') + ;; + push) + files=$(gh api \ + "/repos/${{ github.repository }}/compare/${{ github.event.before }}...${{ github.event.after }}" \ + -q '.files[].filename') + ;; + *) + files="" + ;; + esac + if printf '%s\n' "$files" | grep -q '^tokenspeed-mla/'; then + echo "tokenspeed-mla touched by this run; in-tree source will be installed." + echo "changed=1" >> "$GITHUB_OUTPUT" + else + echo "tokenspeed-mla untouched; keeping pinned PyPI wheel." + echo "changed=0" >> "$GITHUB_OUTPUT" + fi + + unit-test: + needs: scan + if: | + (github.event_name != 'pull_request' || github.event.action != 'closed') && + ( + github.event_name == 'pull_request' || + github.repository == 'lightseekorg/tokenspeed' + ) && + ( + github.event_name != 'pull_request' || + github.event.pull_request.draft == false + ) + strategy: + fail-fast: false + matrix: ${{ fromJson(needs.scan.outputs.matrix) }} + name: ${{ matrix.name }} (${{ matrix.runner }}) + runs-on: ${{ matrix.runner }} + timeout-minutes: 60 + continue-on-error: ${{ matrix.optional }} + env: + # When the scan job determined that this run touched `tokenspeed-mla/`, + # `install_deps.sh` will reinstall the in-tree source on top of the + # pinned PyPI wheel. See the matching detect step in the `scan` job. + INSTALL_TOKENSPEED_MLA_FROM_SOURCE: ${{ needs.scan.outputs.install_tokenspeed_mla_from_source }} + # Authenticate Hugging Face Hub downloads so checkpoint pulls during + # `ts serve` startup are not throttled by the anonymous rate limit. The + # 600s per-file `runtime-1gpu` timeout is otherwise exceeded on cold + # runners (e.g. `amd-mi35x-1gpu-test` loading `openai/gpt-oss-120b` + + # EAGLE3 draft). + HF_TOKEN: ${{ secrets.HF_TOKEN }} + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + path: run-${{ github.run_id }}-${{ github.run_attempt }}-${{ matrix.name }}-${{ matrix.runner }} + + - name: Set work directory + run: | + echo "WORK_DIR=${{ github.workspace }}/run-${{ github.run_id }}-${{ github.run_attempt }}-${{ matrix.name }}-${{ matrix.runner }}" >> "$GITHUB_ENV" + + - name: Install pipeline dependency + # See note in the scan job — pypi.org "Content-Type: Unknown" flake. + run: | + for i in 1 2 3 4 5; do + python3 -m pip install pyyaml && exit 0 + echo "pyyaml install attempt $i failed, retrying in 10s..." + sleep 10 + done + exit 1 + + # When the diff touches `tokenspeed-mla/`, ask install_deps.sh to + # rebuild + install that package from the in-tree source after the + # pinned PyPI wheel has been installed (transitively via + # tokenspeed-kernel). Without this, CI keeps testing whichever + # `tokenspeed-mla` version is pinned in + # `tokenspeed-kernel/python/requirements/cuda-thirdparty.txt`, not + # the code on the PR/branch. When the diff doesn't touch + # `tokenspeed-mla/` we keep using the pinned PyPI wheel. + - name: Install CI task + run: | + cd "${{ env.WORK_DIR }}" + mkdir -p .ci-artifacts + python3 test/ci_system/pipeline.py execute \ + --config "${{ matrix.config }}" \ + --runner "${{ matrix.runner }}" \ + --work-dir "${{ env.WORK_DIR }}" \ + --print-plan \ + --only-stage install \ + --keep-runner-state \ + --result-json .ci-artifacts/result.json + + - name: Execute CI task + run: | + cd "${{ env.WORK_DIR }}" + mkdir -p .ci-artifacts + python3 test/ci_system/pipeline.py execute \ + --config "${{ matrix.config }}" \ + --runner "${{ matrix.runner }}" \ + --work-dir "${{ env.WORK_DIR }}" \ + --print-plan \ + --skip-stage install \ + --reuse-runner-state \ + --result-json .ci-artifacts/result.json + + - name: Upload task result + if: always() + uses: actions/upload-artifact@v4 + with: + name: pr-test-${{ matrix.name }}-${{ matrix.runner }}-${{ github.run_id }}-${{ github.run_attempt }} + path: ${{ env.WORK_DIR }}/.ci-artifacts/result.json + + - name: Cleanup work directory + if: always() + run: rm -rf "${{ env.WORK_DIR }}" + + finish: + if: | + always() && + (github.event_name != 'pull_request' || github.event.action != 'closed') && + ( + github.event_name == 'pull_request' || + github.repository == 'lightseekorg/tokenspeed' + ) + needs: [unit-test] + runs-on: ubuntu-latest + steps: + - name: Finish + run: echo "PR Test finished." diff --git a/.github/workflows/release-docker.yml b/.github/workflows/release-docker.yml new file mode 100644 index 0000000..68bde40 --- /dev/null +++ b/.github/workflows/release-docker.yml @@ -0,0 +1,42 @@ +name: Release Docker Images +on: + push: + branches: + - main + paths: + - "python/tokenspeed/version.py" + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + build: + if: github.repository == 'lightseekorg/tokenspeed' + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + runner_tag: ['cu130-torch-2.11.0', 'cu129-torch-2.11.0'] + steps: + - name: Free disk space + uses: jlumbroso/free-disk-space@main + with: + tool-cache: true + android: true + dotnet: true + haskell: true + large-packages: false + swap-storage: false + + - name: Checkout repository + uses: actions/checkout@v3 + + - name: Build + run: | + docker build . \ + -f docker/Dockerfile \ + --build-arg RUNNER_TAG=${{ matrix.runner_tag }} \ + -t lightseekorg/tokenspeed:${{ matrix.runner_tag }} \ + --no-cache diff --git a/.github/workflows/release-pypi.yml b/.github/workflows/release-pypi.yml new file mode 100644 index 0000000..5ccefbb --- /dev/null +++ b/.github/workflows/release-pypi.yml @@ -0,0 +1,31 @@ +name: Release PyPI +on: + push: + branches: + - main + paths: + - "python/tokenspeed/version.py" + workflow_dispatch: + +jobs: + publish: + if: github.repository == 'lightseekorg/tokenspeed' + 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@v3 + + - 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-tokenspeed-kernel-amd.yml b/.github/workflows/release-tokenspeed-kernel-amd.yml new file mode 100644 index 0000000..0b439af --- /dev/null +++ b/.github/workflows/release-tokenspeed-kernel-amd.yml @@ -0,0 +1,140 @@ +name: Build and Release tokenspeed-kernel-amd + +on: + workflow_dispatch: + inputs: + tag_name: + description: "Git tag to create or update. Defaults to tokenspeed-kernel-amd-v." + required: false + type: string + prerelease: + description: "Mark the GitHub release as a prerelease." + required: false + default: false + type: boolean + +permissions: + contents: read + id-token: write + +concurrency: + group: ${{ github.workflow }} + cancel-in-progress: true + +jobs: + metadata: + name: Resolve release metadata + runs-on: ubuntu-latest + outputs: + version: ${{ steps.metadata.outputs.version }} + tag_name: ${{ steps.metadata.outputs.tag_name }} + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Resolve release metadata + id: metadata + run: | + VERSION=$(python - <<'PY' + import pathlib + import tomllib + + data = tomllib.loads(pathlib.Path("tokenspeed-kernel-amd/pyproject.toml").read_text()) + print(data["project"]["version"]) + PY + ) + TAG_NAME="${{ inputs.tag_name }}" + if [ -z "$TAG_NAME" ]; then + TAG_NAME="tokenspeed-kernel-amd-v${VERSION}" + fi + echo "version=${VERSION}" >> "$GITHUB_OUTPUT" + echo "tag_name=${TAG_NAME}" >> "$GITHUB_OUTPUT" + echo "Release version: ${VERSION}" + echo "Release tag: ${TAG_NAME}" + + build: + name: Build distributions + needs: metadata + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Build distributions + run: | + python -m pip install --upgrade build twine "setuptools<82" wheel + python -m build tokenspeed-kernel-amd --sdist --wheel --no-isolation --outdir dist + python -m twine check dist/* + + - name: Upload distributions + uses: actions/upload-artifact@v4 + with: + name: tokenspeed-kernel-amd-dist + path: dist/* + + release: + name: Create wheelhouse release + needs: [metadata, build] + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Download distributions + uses: actions/download-artifact@v4 + with: + path: dist + pattern: tokenspeed-kernel-amd-* + merge-multiple: true + + - name: Create or update GitHub release + env: + GH_TOKEN: ${{ secrets.LIGHTSEEK_BOT_TOKEN }} + RELEASE_REPOSITORY: lightseekorg/whl + TAG_NAME: ${{ needs.metadata.outputs.tag_name }} + VERSION: ${{ needs.metadata.outputs.version }} + PRERELEASE: ${{ inputs.prerelease }} + SOURCE_SHA: ${{ github.sha }} + run: | + prerelease_args=() + if [ "$PRERELEASE" = "true" ]; then + prerelease_args+=(--prerelease) + fi + + if gh release view "$TAG_NAME" --repo "$RELEASE_REPOSITORY" >/dev/null 2>&1; then + gh release upload "$TAG_NAME" dist/* --repo "$RELEASE_REPOSITORY" --clobber + else + gh release create "$TAG_NAME" dist/* \ + --repo "$RELEASE_REPOSITORY" \ + --title "tokenspeed-kernel-amd ${VERSION}" \ + --notes "tokenspeed-kernel-amd ${VERSION} built from lightseekorg/tokenspeed@${SOURCE_SHA}" \ + "${prerelease_args[@]}" + fi + + publish-pypi: + name: Publish to PyPI + needs: [metadata, build] + runs-on: ubuntu-latest + environment: pypi + steps: + - name: Download distributions + uses: actions/download-artifact@v4 + with: + path: dist + pattern: tokenspeed-kernel-amd-* + merge-multiple: true + + - name: Publish distributions to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 + with: + packages-dir: dist diff --git a/.github/workflows/release-tokenspeed-kernel.yml b/.github/workflows/release-tokenspeed-kernel.yml new file mode 100644 index 0000000..73213ac --- /dev/null +++ b/.github/workflows/release-tokenspeed-kernel.yml @@ -0,0 +1,376 @@ +name: Build and Release tokenspeed-kernel + +on: + workflow_dispatch: + inputs: + tag_name: + description: "Git tag prefix to create or update. Defaults to tokenspeed-kernel-v." + required: false + type: string + version_date: + description: "Version date override for development builds, formatted as YYYYMMDD." + required: false + type: string + cuda_variant: + description: "CUDA variant to build." + required: false + default: all + type: choice + options: + - all + - cu130 + - cu129 + pypi_cuda_variant: + description: "CUDA variant to publish to PyPI." + required: false + default: cu130 + type: choice + options: + - cu130 + - cu129 + prerelease: + description: "Mark the GitHub release as a prerelease." + required: false + default: false + type: boolean + publish_github: + description: "Publish distributions to lightseekorg/whl." + required: false + default: false + type: boolean + publish_pypi: + description: "Publish the selected PyPI CUDA variant." + required: false + default: false + type: boolean + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }} + cancel-in-progress: true + +jobs: + metadata: + name: Resolve release metadata + runs-on: ubuntu-24.04-32core-x64 + outputs: + version: ${{ steps.metadata.outputs.version }} + version_date: ${{ steps.metadata.outputs.version_date }} + tag_name: ${{ steps.metadata.outputs.tag_name }} + steps: + - name: Validate publish inputs + if: inputs.publish_pypi && inputs.cuda_variant != 'all' && inputs.cuda_variant != inputs.pypi_cuda_variant + run: | + echo "pypi_cuda_variant must be included by cuda_variant." >&2 + exit 1 + + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Install build metadata dependencies + run: python -m pip install --upgrade "setuptools<82" wheel packaging + + - name: Resolve release metadata + id: metadata + env: + PUBLISH_PYPI: ${{ inputs.publish_pypi }} + run: | + version_date="${{ inputs.version_date }}" + if [ -z "$version_date" ]; then + version_date="$(date -u +%Y%m%d)" + fi + + version=$( + cd tokenspeed-kernel/python + TOKENSPEED_KERNEL_BACKEND=cuda \ + TOKENSPEED_KERNEL_VERSION_DATE="$version_date" \ + TOKENSPEED_KERNEL_GIT_SHA="$GITHUB_SHA" \ + TOKENSPEED_KERNEL_GIT_BRANCH="$GITHUB_REF_NAME" \ + python setup.py --version + ) + + if [ "$PUBLISH_PYPI" = "true" ] && [[ "$version" == *+* ]]; then + echo "Version ${version} contains a local segment and cannot be published to PyPI; dispatch from a release/* branch." >&2 + exit 1 + fi + + tag_name="${{ inputs.tag_name }}" + if [ -z "$tag_name" ]; then + tag_name="tokenspeed-kernel-v${version}" + fi + + echo "version=${version}" >> "$GITHUB_OUTPUT" + echo "version_date=${version_date}" >> "$GITHUB_OUTPUT" + echo "tag_name=${tag_name}" >> "$GITHUB_OUTPUT" + echo "Release version: ${version}" + echo "Release tag: ${tag_name}" + + build-wheel: + name: Build ${{ matrix.cuda.variant }} ${{ matrix.arch.name }} wheel for Python ${{ matrix.python.version }} + needs: metadata + runs-on: ${{ matrix.arch.runner }} + strategy: + fail-fast: false + matrix: + arch: + - name: x64 + runner: ubuntu-24.04-32core-x64 + manylinux_arch: x86_64 + cuda_repo_arch: x86_64 + - name: arm64 + runner: ubuntu-24.04-32core-arm64 + manylinux_arch: aarch64 + cuda_repo_arch: sbsa + python: + - version: "3.10" + tag: cp310-cp310 + - version: "3.11" + tag: cp311-cp311 + - version: "3.12" + tag: cp312-cp312 + - version: "3.13" + tag: cp313-cp313 + cuda: + - variant: cu130 + package_version: "13-0" + torch_index_url: https://download.pytorch.org/whl/cu130 + - variant: cu129 + package_version: "12-9" + torch_index_url: https://download.pytorch.org/whl/cu129 + steps: + - name: Checkout repository + if: inputs.cuda_variant == 'all' || inputs.cuda_variant == matrix.cuda.variant + uses: actions/checkout@v4 + + - name: Build manylinux CUDA image + if: inputs.cuda_variant == 'all' || inputs.cuda_variant == matrix.cuda.variant + run: | + cat > Dockerfile.tokenspeed-kernel-release <<'DOCKERFILE' + ARG MANYLINUX_ARCH=x86_64 + FROM quay.io/pypa/manylinux_2_28_${MANYLINUX_ARCH} + + ARG CUDA_PKG_VERSION=13-0 + ARG CUDA_REPO_ARCH=x86_64 + + RUN curl -o /etc/yum.repos.d/cuda-rhel8.repo \ + https://developer.download.nvidia.com/compute/cuda/repos/rhel8/${CUDA_REPO_ARCH}/cuda-rhel8.repo \ + && yum install -y \ + cuda-nvcc-${CUDA_PKG_VERSION} \ + cuda-cccl-${CUDA_PKG_VERSION} \ + cuda-cudart-devel-${CUDA_PKG_VERSION} \ + cuda-driver-devel-${CUDA_PKG_VERSION} \ + libcublas-devel-${CUDA_PKG_VERSION} \ + libcurand-devel-${CUDA_PKG_VERSION} \ + libcusolver-devel-${CUDA_PKG_VERSION} \ + libcusparse-devel-${CUDA_PKG_VERSION} \ + gcc-toolset-13-gcc \ + gcc-toolset-13-gcc-c++ \ + ninja-build \ + && yum clean all \ + && rm -rf /var/cache/yum \ + && ln -s /usr/bin/ninja-build /usr/local/bin/ninja + + ENV CC=/opt/rh/gcc-toolset-13/root/usr/bin/gcc + ENV CXX=/opt/rh/gcc-toolset-13/root/usr/bin/g++ + ENV CUDA_HOME=/usr/local/cuda + ENV PATH=${CUDA_HOME}/bin:${PATH} + ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${CUDA_HOME}/lib64/stubs:${LD_LIBRARY_PATH} + ENV LIBRARY_PATH=${CUDA_HOME}/lib64/stubs:${LIBRARY_PATH} + + WORKDIR /workspace + DOCKERFILE + + docker build \ + --build-arg MANYLINUX_ARCH=${{ matrix.arch.manylinux_arch }} \ + --build-arg CUDA_PKG_VERSION=${{ matrix.cuda.package_version }} \ + --build-arg CUDA_REPO_ARCH=${{ matrix.arch.cuda_repo_arch }} \ + -f Dockerfile.tokenspeed-kernel-release \ + -t tokenspeed-kernel-release:${{ matrix.cuda.variant }}-${{ matrix.arch.name }} . + + - name: Build wheel + if: inputs.cuda_variant == 'all' || inputs.cuda_variant == matrix.cuda.variant + run: | + mkdir -p dist + docker run --rm \ + -e PYTHON_TAG=${{ matrix.python.tag }} \ + -e TORCH_INDEX_URL=${{ matrix.cuda.torch_index_url }} \ + -e TOKENSPEED_KERNEL_BACKEND=cuda \ + -e TOKENSPEED_KERNEL_VERSION_DATE=${{ needs.metadata.outputs.version_date }} \ + -e TOKENSPEED_KERNEL_GIT_SHA=$GITHUB_SHA \ + -e TOKENSPEED_KERNEL_GIT_BRANCH=$GITHUB_REF_NAME \ + -e MAX_JOBS=8 \ + -v "$PWD:/workspace" \ + -v "$PWD/dist:/output" \ + tokenspeed-kernel-release:${{ matrix.cuda.variant }}-${{ matrix.arch.name }} \ + bash -lc ' + set -euo pipefail + + PYBIN="/opt/python/${PYTHON_TAG}/bin" + cd /workspace/tokenspeed-kernel/python + + "${PYBIN}/python" -m pip install --upgrade \ + pip "setuptools<82" wheel build twine packaging "apache-tvm-ffi>=0.1.5,<=0.1.11" + "${PYBIN}/python" -m pip install --no-cache-dir "torch==2.11.0" --index-url "${TORCH_INDEX_URL}" + + rm -rf build *.egg-info + PIP_EXTRA_INDEX_URL="${TORCH_INDEX_URL}" \ + "${PYBIN}/python" -m build --wheel --no-isolation --outdir /output + + python_tag="${PYTHON_TAG%%-*}" + abi_tag="${PYTHON_TAG#*-}" + platform_tag="manylinux_2_28_$(uname -m)" + "${PYBIN}/python" -m wheel tags \ + --python-tag "${python_tag}" \ + --abi-tag "${abi_tag}" \ + --platform-tag "${platform_tag}" \ + --remove /output/*.whl + + "${PYBIN}/python" -m twine check /output/*.whl + ' + + - name: Upload wheel + if: inputs.cuda_variant == 'all' || inputs.cuda_variant == matrix.cuda.variant + uses: actions/upload-artifact@v4 + with: + name: tokenspeed-kernel-wheel-${{ matrix.cuda.variant }}-${{ matrix.arch.name }}-py${{ matrix.python.version }} + path: dist/*.whl + + build-sdist: + name: Build source distribution + needs: metadata + runs-on: ubuntu-24.04-32core-x64 + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Build source distribution + run: | + python -m pip install --upgrade \ + build twine "setuptools<82" wheel packaging "apache-tvm-ffi>=0.1.5,<=0.1.11" + cd tokenspeed-kernel/python + TOKENSPEED_KERNEL_BACKEND=cuda \ + TOKENSPEED_KERNEL_VERSION_DATE=${{ needs.metadata.outputs.version_date }} \ + TOKENSPEED_KERNEL_GIT_SHA=$GITHUB_SHA \ + TOKENSPEED_KERNEL_GIT_BRANCH=$GITHUB_REF_NAME \ + python -m build --sdist --no-isolation --outdir ../../dist + python -m twine check ../../dist/*.tar.gz + + - name: Upload source distribution + uses: actions/upload-artifact@v4 + with: + name: tokenspeed-kernel-sdist + path: dist/*.tar.gz + + release: + name: Create wheelhouse release + if: inputs.publish_github + needs: [metadata, build-wheel, build-sdist] + runs-on: ubuntu-24.04-32core-x64 + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Download distributions + uses: actions/download-artifact@v4 + with: + path: dist + pattern: tokenspeed-kernel-* + merge-multiple: false + + - name: Create or update GitHub release + env: + GH_TOKEN: ${{ secrets.LIGHTSEEK_BOT_TOKEN }} + RELEASE_REPOSITORY: lightseekorg/whl + TAG_NAME: ${{ needs.metadata.outputs.tag_name }} + VERSION: ${{ needs.metadata.outputs.version }} + CUDA_VARIANT: ${{ inputs.cuda_variant }} + PRERELEASE: ${{ inputs.prerelease }} + SOURCE_SHA: ${{ github.sha }} + run: | + prerelease_args=() + if [ "$PRERELEASE" = "true" ]; then + prerelease_args+=(--prerelease) + fi + + upload_release() { + local cuda_variant="$1" + local release_tag="${TAG_NAME}-${cuda_variant}" + local release_title="tokenspeed-kernel ${VERSION} ${cuda_variant}" + + shopt -s nullglob + local assets=( + dist/tokenspeed-kernel-wheel-"${cuda_variant}"-*/*.whl + dist/tokenspeed-kernel-sdist/*.tar.gz + ) + shopt -u nullglob + + if [ "${#assets[@]}" -eq 0 ]; then + echo "No distributions found for ${cuda_variant}" >&2 + exit 1 + fi + + if gh release view "$release_tag" --repo "$RELEASE_REPOSITORY" >/dev/null 2>&1; then + gh release upload "$release_tag" "${assets[@]}" \ + --repo "$RELEASE_REPOSITORY" \ + --clobber + else + gh release create "$release_tag" "${assets[@]}" \ + --repo "$RELEASE_REPOSITORY" \ + --title "$release_title" \ + --notes "$release_title built from lightseekorg/tokenspeed@$SOURCE_SHA" \ + "${prerelease_args[@]}" + fi + } + + case "$CUDA_VARIANT" in + all) + upload_release cu130 + upload_release cu129 + ;; + cu130|cu129) + upload_release "$CUDA_VARIANT" + ;; + *) + echo "Unsupported CUDA variant: ${CUDA_VARIANT}" >&2 + exit 1 + ;; + esac + + publish-pypi: + name: Publish ${{ inputs.pypi_cuda_variant }} to PyPI + if: inputs.publish_pypi + needs: [metadata, build-wheel, build-sdist] + runs-on: ubuntu-24.04-32core-x64 + environment: pypi + permissions: + id-token: write + steps: + - name: Download wheels + uses: actions/download-artifact@v4 + with: + path: dist + pattern: tokenspeed-kernel-wheel-${{ inputs.pypi_cuda_variant }}-* + merge-multiple: true + + - name: Download source distribution + uses: actions/download-artifact@v4 + with: + name: tokenspeed-kernel-sdist + path: dist + + - name: Publish distributions to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 + with: + packages-dir: dist diff --git a/.github/workflows/release-tokenspeed-scheduler.yml b/.github/workflows/release-tokenspeed-scheduler.yml new file mode 100644 index 0000000..b6633db --- /dev/null +++ b/.github/workflows/release-tokenspeed-scheduler.yml @@ -0,0 +1,169 @@ +name: Build and Release tokenspeed-scheduler + +on: + workflow_dispatch: + inputs: + tag_name: + description: "Git tag to create or update. Defaults to tokenspeed-scheduler-v." + required: false + type: string + prerelease: + description: "Mark the GitHub release as a prerelease." + required: false + default: false + type: boolean + +permissions: + contents: read + id-token: write + +jobs: + build-wheel: + name: Build wheels (${{ matrix.arch }}) + runs-on: ${{ matrix.runner }} + strategy: + fail-fast: false + matrix: + include: + - arch: x86_64 + runner: ubuntu-latest + build-selector: "cp310-manylinux_x86_64 cp311-manylinux_x86_64 cp312-manylinux_x86_64 cp313-manylinux_x86_64" + - arch: aarch64 + runner: ubuntu-24.04-arm + build-selector: "cp310-manylinux_aarch64 cp311-manylinux_aarch64 cp312-manylinux_aarch64 cp313-manylinux_aarch64" + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Build wheels + uses: pypa/cibuildwheel@v4.1 + env: + CIBW_BEFORE_ALL: dnf install -y openssl-devel + CIBW_BUILD: ${{ matrix.build-selector }} + CIBW_MANYLINUX_AARCH64_IMAGE: manylinux_2_28 + CIBW_MANYLINUX_X86_64_IMAGE: manylinux_2_28 + with: + package-dir: tokenspeed-scheduler + output-dir: dist + + - name: Set up Python for wheel checks + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Check wheels + run: | + python -m pip install --upgrade twine + python -m twine check dist/* + + - name: Upload wheels + uses: actions/upload-artifact@v4 + with: + name: tokenspeed-scheduler-wheels-${{ matrix.arch }} + path: dist/*.whl + + build-sdist: + name: Build source distribution + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Build source distribution + run: | + python -m pip install --upgrade build twine + cd tokenspeed-scheduler + python -m build --sdist --outdir ../dist + python -m twine check ../dist/* + + - name: Upload source distribution + uses: actions/upload-artifact@v4 + with: + name: tokenspeed-scheduler-sdist + path: dist/*.tar.gz + + release: + name: Create wheelhouse release + needs: [build-wheel, build-sdist] + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Resolve release metadata + id: metadata + run: | + VERSION=$(python - <<'PY' + import pathlib + import tomllib + + data = tomllib.loads(pathlib.Path("tokenspeed-scheduler/pyproject.toml").read_text()) + print(data["project"]["version"]) + PY + ) + TAG_NAME="${{ inputs.tag_name }}" + if [ -z "$TAG_NAME" ]; then + TAG_NAME="tokenspeed-scheduler-v${VERSION}" + fi + echo "version=${VERSION}" >> "$GITHUB_OUTPUT" + echo "tag_name=${TAG_NAME}" >> "$GITHUB_OUTPUT" + echo "Release version: ${VERSION}" + echo "Release tag: ${TAG_NAME}" + + - name: Download distributions + uses: actions/download-artifact@v4 + with: + path: dist + pattern: tokenspeed-scheduler-* + merge-multiple: true + + - name: Create or update GitHub release + env: + GH_TOKEN: ${{ secrets.LIGHTSEEK_BOT_TOKEN }} + RELEASE_REPOSITORY: lightseekorg/whl + TAG_NAME: ${{ steps.metadata.outputs.tag_name }} + VERSION: ${{ steps.metadata.outputs.version }} + PRERELEASE: ${{ inputs.prerelease }} + run: | + prerelease_args=() + if [ "$PRERELEASE" = "true" ]; then + prerelease_args+=(--prerelease) + fi + + if gh release view "$TAG_NAME" --repo "$RELEASE_REPOSITORY" >/dev/null 2>&1; then + gh release upload "$TAG_NAME" dist/* --repo "$RELEASE_REPOSITORY" --clobber + else + gh release create "$TAG_NAME" dist/* \ + --repo "$RELEASE_REPOSITORY" \ + --title "tokenspeed-scheduler ${VERSION}" \ + --notes "tokenspeed-scheduler ${VERSION}" \ + "${prerelease_args[@]}" + fi + + publish-pypi: + name: Publish to PyPI + needs: [build-wheel, build-sdist] + runs-on: ubuntu-latest + environment: pypi + steps: + - name: Download distributions + uses: actions/download-artifact@v4 + with: + path: dist + pattern: tokenspeed-scheduler-* + merge-multiple: true + + - name: Publish distributions to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 + with: + packages-dir: dist diff --git a/.github/workflows/retry-failed-approved-prs.yml b/.github/workflows/retry-failed-approved-prs.yml new file mode 100644 index 0000000..51750b1 --- /dev/null +++ b/.github/workflows/retry-failed-approved-prs.yml @@ -0,0 +1,350 @@ +name: Retry Failed Approved PR CI + +on: + schedule: + - cron: "47 * * * *" + workflow_dispatch: + inputs: + dry_run: + description: "Log eligible reruns without starting them" + required: false + default: false + type: boolean + +permissions: + actions: write + contents: read + pull-requests: read + +concurrency: + group: ${{ github.workflow }} + cancel-in-progress: false + +jobs: + rerun-failed: + name: Rerun failed approved-PR jobs + if: github.repository == 'lightseekorg/tokenspeed' + runs-on: ubuntu-latest + timeout-minutes: 15 + steps: + - name: Rerun failed vendor CI jobs + uses: actions/github-script@v7 + with: + script: | + const workflows = [ + "pr-test-nvidia.yml", + "pr-test-amd.yml", + "pr-test-nvidia-arm.yml", + ]; + const maxAttempts = 3; + const maxRunAgeMs = 30 * 24 * 60 * 60 * 1000; + const dryRun = + context.eventName === "workflow_dispatch" && + String(context.payload.inputs?.dry_run) === "true"; + const { owner, repo } = context.repo; + const summaryRows = []; + + function addSummaryRow(pull, workflowId, run, result) { + summaryRows.push([ + `#${pull.number}`, + workflowId || "—", + run + ? `${run.id}` + : "—", + result, + ]); + } + + async function getPullRequest(number) { + const query = ` + query($owner: String!, $repo: String!, $number: Int!) { + repository(owner: $owner, name: $repo) { + pullRequest(number: $number) { + state + isDraft + headRefOid + reviewDecision + author { + login + } + } + } + } + `; + const result = await github.graphql(query, { + owner, + repo, + number, + }); + return result.repository.pullRequest; + } + + async function hasFallbackApproval(number, authorLogin) { + const reviews = await github.paginate( + github.rest.pulls.listReviews, + { + owner, + repo, + pull_number: number, + per_page: 100, + }, + ); + const latestDecisiveReview = new Map(); + + for (const review of reviews) { + if ( + !review.user || + review.user.login === authorLogin || + !["APPROVED", "CHANGES_REQUESTED"].includes(review.state) + ) { + continue; + } + const previous = latestDecisiveReview.get(review.user.id); + if ( + !previous || + Date.parse(review.submitted_at) >= + Date.parse(previous.submitted_at) + ) { + latestDecisiveReview.set(review.user.id, review); + } + } + + const decisions = [...latestDecisiveReview.values()].map( + (review) => review.state, + ); + return ( + decisions.includes("APPROVED") && + !decisions.includes("CHANGES_REQUESTED") + ); + } + + async function isEligible(number, pullRequest) { + if ( + !pullRequest || + pullRequest.state !== "OPEN" || + pullRequest.isDraft + ) { + return false; + } + if (pullRequest.reviewDecision === "APPROVED") { + return true; + } + if (pullRequest.reviewDecision !== null) { + core.info( + "PR #" + number + ": review decision is " + + pullRequest.reviewDecision + ".", + ); + return false; + } + + const approved = await hasFallbackApproval( + number, + pullRequest.author?.login, + ); + if (!approved) { + core.info("PR #" + number + ": no effective approval."); + } + return approved; + } + + try { + const pulls = await github.paginate(github.rest.pulls.list, { + owner, + repo, + state: "open", + per_page: 100, + }); + + core.info( + "Scanning " + pulls.length + " open pull request(s)" + + (dryRun ? " in dry-run mode." : "."), + ); + + for (const pull of pulls) { + try { + const pullRequest = await getPullRequest(pull.number); + if (!(await isEligible(pull.number, pullRequest))) { + continue; + } + + const selectedRuns = []; + for (const workflowId of workflows) { + let runs; + try { + runs = await github.paginate( + github.rest.actions.listWorkflowRuns, + { + owner, + repo, + workflow_id: workflowId, + event: "pull_request", + head_sha: pullRequest.headRefOid, + per_page: 100, + }, + ); + } catch (error) { + core.warning( + "PR #" + pull.number + ": unable to list runs for " + + workflowId + ": " + error.message, + ); + addSummaryRow( + pull, + workflowId, + null, + "Failed to inspect", + ); + continue; + } + + const run = runs + .filter( + (candidate) => + candidate.head_sha === pullRequest.headRefOid && + candidate.event === "pull_request", + ) + .sort( + (left, right) => + Date.parse(right.created_at) - + Date.parse(left.created_at), + )[0]; + if (run) { + selectedRuns.push({ workflowId, run }); + } else { + core.info( + "PR #" + pull.number + ": " + workflowId + + " has no run for the current head; not applicable.", + ); + } + } + + if (selectedRuns.length === 0) { + core.info( + "PR #" + pull.number + ": no vendor CI runs found.", + ); + continue; + } + + const failedRuns = selectedRuns.filter( + ({ run }) => + run.status === "completed" && + run.conclusion === "failure", + ); + if (failedRuns.length === 0) { + core.info( + "PR #" + pull.number + ": vendor CI has no failures.", + ); + continue; + } + + // Keep the approval and head SHA checks close to the + // mutations. A push can still race this small window, just as + // it can in the latest-main retry workflow. + const currentPullRequest = await getPullRequest(pull.number); + if ( + currentPullRequest?.headRefOid !== pullRequest.headRefOid || + !(await isEligible(pull.number, currentPullRequest)) + ) { + core.info( + "PR #" + pull.number + + ": eligibility or head changed; skipping reruns.", + ); + continue; + } + + for (const { workflowId, run } of failedRuns) { + const attempt = run.run_attempt || 1; + if (attempt >= maxAttempts) { + core.warning( + "PR #" + pull.number + ": " + workflowId + " run " + + run.id + " failed on attempt " + attempt + + "; retry limit reached.", + ); + addSummaryRow( + pull, + workflowId, + run, + "Retry limit reached", + ); + continue; + } + if ( + Date.now() - Date.parse(run.created_at) > + maxRunAgeMs + ) { + core.warning( + "PR #" + pull.number + ": " + workflowId + " run " + + run.id + " is older than 30 days; skipping.", + ); + addSummaryRow( + pull, + workflowId, + run, + "Run older than 30 days", + ); + continue; + } + + if (dryRun) { + core.notice( + "PR #" + pull.number + ": would rerun failed jobs for " + + workflowId + " run " + run.id + " after attempt " + + attempt + ".", + ); + addSummaryRow(pull, workflowId, run, "Dry run"); + continue; + } + + try { + await github.request( + "POST /repos/{owner}/{repo}/actions/runs/{run_id}/rerun-failed-jobs", + { owner, repo, run_id: run.id }, + ); + core.notice( + "PR #" + pull.number + ": rerunning failed jobs for " + + workflowId + " run " + run.id + " after attempt " + + attempt + ".", + ); + addSummaryRow( + pull, + workflowId, + run, + "Rerun started", + ); + } catch (error) { + core.warning( + "PR #" + pull.number + ": unable to rerun failed jobs " + + "for " + workflowId + " run " + run.id + ": " + + error.message, + ); + addSummaryRow( + pull, + workflowId, + run, + "Failed to start", + ); + } + } + } catch (error) { + core.warning( + "PR #" + pull.number + ": scan failed: " + error.message, + ); + addSummaryRow(pull, null, null, "Failed to inspect"); + } + } + } finally { + core.summary.addHeading("Approved PR retry summary"); + if (summaryRows.length > 0) { + core.summary.addTable([ + [ + { data: "PR", header: true }, + { data: "Workflow", header: true }, + { data: "Run", header: true }, + { data: "Result", header: true }, + ], + ...summaryRows, + ]); + } else { + core.summary.addRaw("No reruns were eligible."); + } + await core.summary.write(); + } diff --git a/.github/workflows/retry-failed-latest-main.yml b/.github/workflows/retry-failed-latest-main.yml new file mode 100644 index 0000000..e272d46 --- /dev/null +++ b/.github/workflows/retry-failed-latest-main.yml @@ -0,0 +1,193 @@ +name: Retry Failed Latest Main CI + +on: + schedule: + - cron: "17 * * * *" + workflow_dispatch: + +permissions: + actions: write + contents: read + +concurrency: + group: ${{ github.workflow }} + cancel-in-progress: false + +jobs: + rerun-failed: + name: Rerun failed latest-main jobs + if: github.repository == 'lightseekorg/tokenspeed' + runs-on: ubuntu-latest + timeout-minutes: 10 + steps: + - name: Rerun failed vendor CI jobs + uses: actions/github-script@v7 + with: + script: | + const workflows = [ + "pr-test-nvidia.yml", + "pr-test-amd.yml", + "pr-test-nvidia-arm.yml", + ]; + const maxAttempts = 3; + const { owner, repo } = context.repo; + const summaryRows = []; + let stoppedEarly = false; + + function addSummaryRow(workflowId, run, result) { + summaryRows.push([ + workflowId, + run + ? `${run.id}` + : "—", + result, + ]); + } + + try { + const { data: repository } = await github.rest.repos.get({ + owner, + repo, + }); + const branch = repository.default_branch; + const { data: commit } = await github.rest.repos.getCommit({ + owner, + repo, + ref: branch, + }); + const latestSha = commit.sha; + + core.info( + "Scanning " + owner + "/" + repo + "@" + latestSha + + " on " + branch, + ); + + for (const workflowId of workflows) { + let runs; + try { + runs = await github.paginate( + github.rest.actions.listWorkflowRuns, + { + owner, + repo, + workflow_id: workflowId, + branch, + event: "push", + head_sha: latestSha, + per_page: 100, + }, + ); + } catch (error) { + core.warning( + "Unable to list runs for " + workflowId + ": " + error.message, + ); + addSummaryRow(workflowId, null, "Failed to inspect"); + continue; + } + + const run = runs + .filter( + (candidate) => + candidate.head_sha === latestSha && + candidate.head_branch === branch && + candidate.event === "push", + ) + .sort( + (left, right) => + Date.parse(right.created_at) - Date.parse(left.created_at), + )[0]; + + if (!run) { + core.info(workflowId + ": no run for latest main; skipping."); + continue; + } + if (run.status !== "completed") { + core.info(workflowId + ": run " + run.id + " is " + run.status + "."); + continue; + } + if (run.conclusion !== "failure") { + core.info( + workflowId + ": run " + run.id + " concluded " + + run.conclusion + "; skipping.", + ); + continue; + } + + const attempt = run.run_attempt || 1; + if (attempt >= maxAttempts) { + core.warning( + workflowId + ": run " + run.id + " failed on attempt " + + attempt + "; retry limit reached.", + ); + addSummaryRow(workflowId, run, "Retry limit reached"); + continue; + } + + // Best-effort guard against main advancing during this scan. + let currentSha; + try { + const { data: currentCommit } = + await github.rest.repos.getCommit({ + owner, + repo, + ref: branch, + }); + currentSha = currentCommit.sha; + } catch (error) { + core.warning( + "Unable to recheck latest main: " + error.message + + "; stopping this scan.", + ); + stoppedEarly = true; + break; + } + if (currentSha !== latestSha) { + core.info( + "Main advanced from " + latestSha + " to " + currentSha + + "; stopping this scan.", + ); + stoppedEarly = true; + break; + } + + try { + await github.request( + "POST /repos/{owner}/{repo}/actions/runs/{run_id}/rerun-failed-jobs", + { owner, repo, run_id: run.id }, + ); + core.notice( + workflowId + ": rerunning failed jobs for run " + run.id + + " after attempt " + attempt + ".", + ); + addSummaryRow(workflowId, run, "Rerun started"); + } catch (error) { + core.warning( + workflowId + ": unable to rerun failed jobs for run " + + run.id + ": " + error.message, + ); + addSummaryRow(workflowId, run, "Failed to start"); + } + } + } finally { + core.summary.addHeading("Latest main retry summary"); + if (summaryRows.length > 0) { + core.summary.addTable([ + [ + { data: "Workflow", header: true }, + { data: "Run", header: true }, + { data: "Result", header: true }, + ], + ...summaryRows, + ]); + } else { + core.summary.addRaw("No reruns were eligible."); + } + if (stoppedEarly) { + core.summary.addBreak(); + core.summary.addRaw( + "Scan stopped early because main advanced or could not " + + "be re-verified.", + ); + } + await core.summary.write(); + } diff --git a/.github/workflows/scheduler-cpp-test.yml b/.github/workflows/scheduler-cpp-test.yml new file mode 100644 index 0000000..a6bcfab --- /dev/null +++ b/.github/workflows/scheduler-cpp-test.yml @@ -0,0 +1,82 @@ +name: Scheduler C++ Test + +on: + push: + branches: [main] + paths: + - "tokenspeed-scheduler/**" + - ".github/workflows/scheduler-cpp-test.yml" + pull_request: + branches: [main] + paths: + - "tokenspeed-scheduler/**" + - ".github/workflows/scheduler-cpp-test.yml" + workflow_dispatch: + +concurrency: + group: scheduler-cpp-test-${{ github.ref }} + cancel-in-progress: true + +jobs: + test: + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.12" + + - name: Install dependencies + run: | + sudo apt-get update + sudo apt-get install -y cmake ninja-build + python -m pip install "nanobind>=2.13.0" tokenspeed-spdlog==1.15.1 + + - name: Configure scheduler C++ tests + run: | + cmake -S tokenspeed-scheduler -B tokenspeed-scheduler/build \ + -G Ninja \ + -DTOKENSPEED_SCHEDULER_BUILD_TESTS=ON \ + -DTOKENSPEED_SCHEDULER_BUILD_PYTHON=OFF + + - name: Build scheduler C++ tests + run: cmake --build tokenspeed-scheduler/build --target tokenspeed_scheduler_tests + + - name: Run scheduler C++ tests + run: ./tokenspeed-scheduler/build/tokenspeed_scheduler_tests + + test-flat: + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.12" + + - name: Install dependencies + run: | + sudo apt-get update + sudo apt-get install -y cmake ninja-build + python -m pip install "nanobind>=2.13.0" tokenspeed-spdlog==1.15.1 + + - name: Configure scheduler C++ tests (flat KV cache) + run: | + cmake -S tokenspeed-scheduler -B tokenspeed-scheduler/build-flat \ + -G Ninja \ + -DTOKENSPEED_SCHEDULER_BUILD_TESTS=ON \ + -DTOKENSPEED_SCHEDULER_BUILD_PYTHON=OFF \ + -DTOKENSPEED_FLAT_KVCACHE=ON + + - name: Build scheduler C++ tests (flat KV cache) + run: cmake --build tokenspeed-scheduler/build-flat --target tokenspeed_scheduler_tests + + - name: Run scheduler C++ tests (flat KV cache) + run: | + FILTER="-$(grep -v '^#' tokenspeed-scheduler/tests/cpp/flat_known_radix_failures.txt | paste -sd: -)" + ./tokenspeed-scheduler/build-flat/tokenspeed_scheduler_tests --gtest_filter="$FILTER" diff --git a/.github/workflows/scheduler-python-test.yml b/.github/workflows/scheduler-python-test.yml new file mode 100644 index 0000000..2bab2c9 --- /dev/null +++ b/.github/workflows/scheduler-python-test.yml @@ -0,0 +1,39 @@ +name: Scheduler Python Test + +on: + push: + branches: [main] + paths: + - "tokenspeed-scheduler/**" + - ".github/workflows/scheduler-python-test.yml" + pull_request: + branches: [main] + paths: + - "tokenspeed-scheduler/**" + - ".github/workflows/scheduler-python-test.yml" + workflow_dispatch: + +concurrency: + group: scheduler-python-test-${{ github.ref }} + cancel-in-progress: true + +jobs: + test: + runs-on: ubuntu-latest + steps: + - name: Checkout code + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.12" + + - name: Install dependencies + run: | + sudo apt-get update + sudo apt-get install -y cmake ninja-build + python -m pip install -e 'tokenspeed-scheduler[test]' + + - name: Run scheduler Python tests + run: python -m pytest tokenspeed-scheduler/python/tests diff --git a/.github/workflows/stale.yml b/.github/workflows/stale.yml new file mode 100644 index 0000000..7fac17a --- /dev/null +++ b/.github/workflows/stale.yml @@ -0,0 +1,37 @@ +name: Close Inactive Issues and PRs + +on: + schedule: + - cron: '0 0 * * *' + workflow_dispatch: + +permissions: + actions: write + issues: write + pull-requests: write + +jobs: + stale: + if: github.repository == 'lightseekorg/tokenspeed' + runs-on: ubuntu-latest + steps: + - uses: actions/stale@v10 + with: + days-before-stale: 14 + days-before-close: 3 + stale-issue-label: inactive + stale-issue-message: > + This issue has been inactive for 14 days and is marked as stale. + It will be closed in 3 days if there is no further activity. + close-issue-message: > + This issue has been automatically closed due to inactivity. + Feel free to reopen it if it is still relevant. + exempt-issue-labels: 'good first issue' + stale-pr-label: inactive + stale-pr-message: > + This PR has been inactive for 14 days and is marked as stale. + It will be closed in 3 days if there is no further activity. + close-pr-message: > + This PR has been automatically closed due to inactivity. + Feel free to reopen it if it is still relevant. + operations-per-run: 200 diff --git a/.github/workflows/update-tokenspeed-kernel-flashinfer.yml b/.github/workflows/update-tokenspeed-kernel-flashinfer.yml new file mode 100644 index 0000000..28c8d5b --- /dev/null +++ b/.github/workflows/update-tokenspeed-kernel-flashinfer.yml @@ -0,0 +1,104 @@ +name: Update tokenspeed-kernel FlashInfer + +on: + workflow_dispatch: + inputs: + flashinfer_version: + description: "FlashInfer version to use, for example 0.6.11 or v0.6.11." + required: true + type: string + +permissions: + contents: write + pull-requests: write + +concurrency: + group: ${{ github.workflow }}-${{ inputs.flashinfer_version }} + cancel-in-progress: true + +jobs: + update: + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@v4 + with: + ref: main + + - name: Configure git author + run: | + git config user.name "lightseek-bot" + git config user.email "243258330+lightseek-bot@users.noreply.github.com" + + - name: Update FlashInfer dependency + id: update + run: | + raw_version="${{ inputs.flashinfer_version }}" + version="${raw_version#v}" + + if ! [[ "$version" =~ ^[0-9]+(\.[0-9]+){1,3}([a-zA-Z0-9.+-]+)?$ ]]; then + echo "Invalid FlashInfer version: $raw_version" >&2 + exit 1 + fi + + python - "$version" <<'PY' + import re + import sys + from pathlib import Path + + version = sys.argv[1] + path = Path("tokenspeed-kernel/python/requirements/cuda.txt") + text = path.read_text() + replacements = [ + (r"flashinfer-python==[^\n]+", f"flashinfer-python=={version}"), + (r"flashinfer-cubin==[^\n]+", f"flashinfer-cubin=={version}"), + ] + for pattern, replacement in replacements: + text = re.sub(pattern, replacement, text) + path.write_text(text) + PY + + echo "version=$version" >> "$GITHUB_OUTPUT" + + - name: Create pull request + env: + GH_TOKEN: ${{ secrets.LIGHTSEEK_BOT_TOKEN }} + VERSION: ${{ steps.update.outputs.version }} + run: | + if [ -z "$GH_TOKEN" ]; then + echo "LIGHTSEEK_BOT_TOKEN is required because this repository does not allow GITHUB_TOKEN to create pull requests." >&2 + exit 1 + fi + + if git diff --quiet; then + echo "FlashInfer is already at $VERSION." + exit 0 + fi + + branch="bot/update-tokenspeed-kernel-flashinfer-${VERSION}" + git checkout -b "$branch" + git add tokenspeed-kernel/python/requirements/cuda.txt + git commit -s -m "Update tokenspeed-kernel FlashInfer to ${VERSION}" + git push --force origin "HEAD:refs/heads/${branch}" + + cat > pr-body.md </dev/null 2>&1; then + gh pr edit "$branch" \ + --repo "$GITHUB_REPOSITORY" \ + --title "Update tokenspeed-kernel FlashInfer to ${VERSION}" \ + --body-file pr-body.md + else + gh pr create \ + --repo "$GITHUB_REPOSITORY" \ + --base main \ + --head "$branch" \ + --title "Update tokenspeed-kernel FlashInfer to ${VERSION}" \ + --body-file pr-body.md + fi diff --git a/.github/workflows/update-tokenspeed-kernel-mla.yml b/.github/workflows/update-tokenspeed-kernel-mla.yml new file mode 100644 index 0000000..a7d86cb --- /dev/null +++ b/.github/workflows/update-tokenspeed-kernel-mla.yml @@ -0,0 +1,106 @@ +name: Update tokenspeed-kernel MLA + +on: + workflow_dispatch: + inputs: + mla_version: + description: "tokenspeed-mla version to use, for example 0.1.2 or v0.1.2." + required: true + type: string + +permissions: + contents: write + pull-requests: write + +concurrency: + group: ${{ github.workflow }}-${{ inputs.mla_version }} + cancel-in-progress: true + +jobs: + update: + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@v4 + with: + ref: main + + - name: Configure git author + run: | + git config user.name "lightseek-bot" + git config user.email "243258330+lightseek-bot@users.noreply.github.com" + + - name: Update tokenspeed-mla dependency + id: update + run: | + raw_version="${{ inputs.mla_version }}" + version="${raw_version#v}" + + if ! [[ "$version" =~ ^[0-9]+(\.[0-9]+)*((a|b|rc)[0-9]+)?(\.post[0-9]+)?(\.dev[0-9]+)?(\+[A-Za-z0-9.]+)?$ ]]; then + echo "Invalid tokenspeed-mla version: $raw_version" >&2 + exit 1 + fi + + python - "$version" <<'PY' + import re + import sys + from pathlib import Path + + version = sys.argv[1] + path = Path("tokenspeed-kernel/python/requirements/cuda-thirdparty.txt") + text = path.read_text() + text, count = re.subn( + r"(?m)^tokenspeed-mla==[^\n]+$", + f"tokenspeed-mla=={version}", + text, + count=1, + ) + if count != 1: + raise SystemExit("tokenspeed-mla dependency not found") + path.write_text(text) + PY + + echo "version=$version" >> "$GITHUB_OUTPUT" + + - name: Create pull request + env: + GH_TOKEN: ${{ secrets.LIGHTSEEK_BOT_TOKEN }} + VERSION: ${{ steps.update.outputs.version }} + run: | + if [ -z "$GH_TOKEN" ]; then + echo "LIGHTSEEK_BOT_TOKEN is required because this repository does not allow GITHUB_TOKEN to create pull requests." >&2 + exit 1 + fi + + if git diff --quiet; then + echo "tokenspeed-kernel already uses tokenspeed-mla $VERSION." + exit 0 + fi + + branch="bot/update-tokenspeed-kernel-mla-${VERSION//+/-}" + git checkout -b "$branch" + git add tokenspeed-kernel/python/requirements/cuda-thirdparty.txt + git commit -s -m "Update tokenspeed-kernel MLA to ${VERSION}" + git push --force origin "HEAD:refs/heads/${branch}" + + cat > pr-body.md </dev/null 2>&1; then + gh pr edit "$branch" \ + --repo "$GITHUB_REPOSITORY" \ + --title "Update tokenspeed-kernel MLA to ${VERSION}" \ + --body-file pr-body.md + else + gh pr create \ + --repo "$GITHUB_REPOSITORY" \ + --base main \ + --head "$branch" \ + --title "Update tokenspeed-kernel MLA to ${VERSION}" \ + --body-file pr-body.md + fi diff --git a/.github/workflows/update-tokenspeed-mla-version.yml b/.github/workflows/update-tokenspeed-mla-version.yml new file mode 100644 index 0000000..a865319 --- /dev/null +++ b/.github/workflows/update-tokenspeed-mla-version.yml @@ -0,0 +1,117 @@ +name: Update tokenspeed-mla version + +on: + workflow_dispatch: + inputs: + version: + description: "tokenspeed-mla version to write to pyproject.toml, for example 0.1.2." + required: true + type: string + +permissions: + contents: write + pull-requests: write + +concurrency: + group: ${{ github.workflow }}-${{ inputs.version }} + cancel-in-progress: true + +jobs: + update: + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@v4 + with: + ref: main + + - name: Configure git author + run: | + git config user.name "lightseek-bot" + git config user.email "243258330+lightseek-bot@users.noreply.github.com" + + - name: Update tokenspeed-mla version + id: update + run: | + version="${{ inputs.version }}" + + if ! [[ "$version" =~ ^[0-9]+(\.[0-9]+)*((a|b|rc)[0-9]+)?(\.post[0-9]+)?(\.dev[0-9]+)?(\+[A-Za-z0-9.]+)?$ ]]; then + echo "Invalid tokenspeed-mla version: $version" >&2 + exit 1 + fi + + python - "$version" <<'PY' + import re + import sys + import tomllib + from pathlib import Path + + version = sys.argv[1] + path = Path("tokenspeed-mla/pyproject.toml") + text = path.read_text() + data = tomllib.loads(text) + if data["project"]["name"] != "tokenspeed-mla": + raise SystemExit("tokenspeed-mla/pyproject.toml does not describe tokenspeed-mla") + + project_match = re.search(r"(?ms)^\[project\]\n(?P.*?)(?=^\[|\Z)", text) + if project_match is None: + raise SystemExit("missing [project] table") + + body = project_match.group("body") + updated_body, count = re.subn( + r'(?m)^version = "[^"]+"$', + f'version = "{version}"', + body, + count=1, + ) + if count != 1: + raise SystemExit("missing project.version") + + text = text[: project_match.start("body")] + updated_body + text[project_match.end("body") :] + path.write_text(text) + PY + + echo "version=$version" >> "$GITHUB_OUTPUT" + + - name: Create pull request + env: + GH_TOKEN: ${{ secrets.LIGHTSEEK_BOT_TOKEN }} + VERSION: ${{ steps.update.outputs.version }} + run: | + if [ -z "$GH_TOKEN" ]; then + echo "LIGHTSEEK_BOT_TOKEN is required because this repository does not allow GITHUB_TOKEN to create pull requests." >&2 + exit 1 + fi + + if git diff --quiet; then + echo "tokenspeed-mla is already at $VERSION." + exit 0 + fi + + branch="bot/update-tokenspeed-mla-version-${VERSION//+/-}" + git checkout -b "$branch" + git add tokenspeed-mla/pyproject.toml + git commit -s -m "Update tokenspeed-mla to ${VERSION}" + git push --force origin "HEAD:refs/heads/${branch}" + + cat > pr-body.md </dev/null 2>&1; then + gh pr edit "$branch" \ + --repo "$GITHUB_REPOSITORY" \ + --title "Update tokenspeed-mla to ${VERSION}" \ + --body-file pr-body.md + else + gh pr create \ + --repo "$GITHUB_REPOSITORY" \ + --base main \ + --head "$branch" \ + --title "Update tokenspeed-mla to ${VERSION}" \ + --body-file pr-body.md + fi diff --git a/.github/workflows/update-tokenspeed-smg-version.yml b/.github/workflows/update-tokenspeed-smg-version.yml new file mode 100644 index 0000000..16779d1 --- /dev/null +++ b/.github/workflows/update-tokenspeed-smg-version.yml @@ -0,0 +1,137 @@ +name: Update tokenspeed-smg versions + +on: + workflow_dispatch: + inputs: + smg_version: + description: "tokenspeed-smg version pin, for example 1.4.1.post20260519." + required: true + type: string + smg_grpc_proto_version: + description: "tokenspeed-smg-grpc-proto version pin, for example 0.4.7.post20260519." + required: true + type: string + smg_grpc_servicer_version: + description: "tokenspeed-smg-grpc-servicer version pin, for example 0.5.3.post20260519." + required: true + type: string + +permissions: + contents: write + pull-requests: write + +concurrency: + group: ${{ github.workflow }}-${{ inputs.smg_version }}-${{ inputs.smg_grpc_proto_version }}-${{ inputs.smg_grpc_servicer_version }} + cancel-in-progress: true + +jobs: + update: + runs-on: ubuntu-latest + steps: + - name: Checkout repository + uses: actions/checkout@v4 + with: + ref: main + + - name: Configure git author + run: | + git config user.name "lightseek-bot" + git config user.email "243258330+lightseek-bot@users.noreply.github.com" + + - name: Update tokenspeed-smg* pins + id: update + env: + SMG_VERSION: ${{ inputs.smg_version }} + SMG_GRPC_PROTO_VERSION: ${{ inputs.smg_grpc_proto_version }} + SMG_GRPC_SERVICER_VERSION: ${{ inputs.smg_grpc_servicer_version }} + run: | + set -euo pipefail + + pep440='^[0-9]+(\.[0-9]+)*((a|b|rc)[0-9]+)?(\.post[0-9]+)?(\.dev[0-9]+)?(\+[A-Za-z0-9.]+)?$' + for name_var in \ + "tokenspeed-smg:SMG_VERSION" \ + "tokenspeed-smg-grpc-proto:SMG_GRPC_PROTO_VERSION" \ + "tokenspeed-smg-grpc-servicer:SMG_GRPC_SERVICER_VERSION"; do + name="${name_var%:*}" + var="${name_var##*:}" + value="${!var}" + if ! [[ "$value" =~ $pep440 ]]; then + echo "Invalid $name version: $value" >&2 + exit 1 + fi + done + + python - <<'PY' + import os + import re + from pathlib import Path + + targets = { + "tokenspeed-smg": os.environ["SMG_VERSION"], + "tokenspeed-smg-grpc-proto": os.environ["SMG_GRPC_PROTO_VERSION"], + "tokenspeed-smg-grpc-servicer": os.environ["SMG_GRPC_SERVICER_VERSION"], + } + path = Path("python/pyproject.toml") + text = path.read_text() + # Update the longer names first so the prefix match doesn't trip + # on tokenspeed-smg matching tokenspeed-smg-grpc-* lines. + for name in sorted(targets, key=len, reverse=True): + version = targets[name] + pattern = rf'(?m)^(?P\s*")(?P{re.escape(name)})==[^"]+(?P",?)\s*$' + text, count = re.subn(pattern, lambda m: f'{m["lead"]}{m["name"]}=={version}{m["tail"]}', text, count=1) + if count != 1: + raise SystemExit(f"pin for {name} not found in python/pyproject.toml") + path.write_text(text) + PY + + echo "smg_version=${SMG_VERSION}" >> "$GITHUB_OUTPUT" + echo "smg_grpc_proto_version=${SMG_GRPC_PROTO_VERSION}" >> "$GITHUB_OUTPUT" + echo "smg_grpc_servicer_version=${SMG_GRPC_SERVICER_VERSION}" >> "$GITHUB_OUTPUT" + + - name: Create pull request + env: + GH_TOKEN: ${{ secrets.LIGHTSEEK_BOT_TOKEN }} + SMG_VERSION: ${{ steps.update.outputs.smg_version }} + SMG_GRPC_PROTO_VERSION: ${{ steps.update.outputs.smg_grpc_proto_version }} + SMG_GRPC_SERVICER_VERSION: ${{ steps.update.outputs.smg_grpc_servicer_version }} + run: | + if [ -z "$GH_TOKEN" ]; then + echo "LIGHTSEEK_BOT_TOKEN is required because this repository does not allow GITHUB_TOKEN to create pull requests." >&2 + exit 1 + fi + + if git diff --quiet; then + echo "tokenspeed-smg* pins already match the requested versions." + exit 0 + fi + + branch="bot/update-tokenspeed-smg-${SMG_VERSION//+/-}" + git checkout -b "$branch" + git add python/pyproject.toml + git commit -s -m "Update tokenspeed-smg* pins to ${SMG_VERSION} / ${SMG_GRPC_PROTO_VERSION} / ${SMG_GRPC_SERVICER_VERSION}" + git push --force origin "HEAD:refs/heads/${branch}" + + cat > pr-body.md </dev/null 2>&1; then + gh pr edit "$branch" \ + --repo "$GITHUB_REPOSITORY" \ + --title "$title" \ + --body-file pr-body.md + else + gh pr create \ + --repo "$GITHUB_REPOSITORY" \ + --base main \ + --head "$branch" \ + --title "$title" \ + --body-file pr-body.md + fi diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..0b1565a --- /dev/null +++ b/.gitignore @@ -0,0 +1,198 @@ +# VSCode +.vscode + +# ============================================================ +# Python +# ============================================================ + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[codz] +*$py.class + +# 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 +*.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/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Sphinx documentation +docs/_build/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# .python-version + +# pdm +.pdm-python +.pdm-build/ + +# pixi +.pixi + +# PEP 582 +__pypackages__/ + +# Celery +celerybeat-schedule +celerybeat.pid + +# SageMath +*.sage.py + +# Environments +.python-version +.env +.envrc +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder +.spyderproject +.spyproject + +# Rope +.ropeproject + +# mkdocs +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre +.pyre/ + +# pytype +.pytype/ + +# Cython debug symbols +cython_debug/ + +# Ruff +.ruff_cache/ + +# PyPI +.pypirc + +# ============================================================ +# C++ +# ============================================================ + +# Prerequisites +*.d + +# Compiled Object files +*.slo +*.lo +*.o +*.obj + +# Precompiled Headers +*.gch +*.pch + +# Linker files +*.ilk + +# Debugger files +*.pdb + +# Compiled Dynamic libraries +*.so +*.dylib +*.dll +*.so.* + +# Fortran module files +*.mod +*.smod + +# Compiled Static libraries +*.lai +*.la +*.a +*.lib + +# Executables +*.exe +*.out +*.app + +# Build directories +Build/ +build-*/ +build_*/ + +# CMake generated files +CMakeFiles/ +CMakeCache.txt +cmake_install.cmake +install_manifest.txt +compile_commands.json + +# vcpkg +vcpkg_installed/ + +# Debug information +*.dwo + +# Test output & cache +Testing/ +.claude/* +.serena/* + +AGENTS.local.md diff --git a/.isort.cfg b/.isort.cfg new file mode 100644 index 0000000..15a1ba0 --- /dev/null +++ b/.isort.cfg @@ -0,0 +1,3 @@ +[settings] +profile=black +known_first_party=tokenspeed diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..47d2d06 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,37 @@ +default_stages: [pre-commit, pre-push, manual] + +default_language_version: + python: python3.12 + +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v6.0.0 + hooks: + - id: check-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: 8.0.1 + hooks: + - id: isort + - repo: https://github.com/psf/black + rev: 26.3.1 + hooks: + - id: black-jupyter + - repo: https://github.com/kynan/nbstripout + rev: 0.9.1 + hooks: + - id: nbstripout + args: + - '--keep-output' + - '--extra-keys=metadata.kernelspec metadata.language_info.version' diff --git a/.skills/bisect-triton-release.md b/.skills/bisect-triton-release.md new file mode 100644 index 0000000..6c3ed24 --- /dev/null +++ b/.skills/bisect-triton-release.md @@ -0,0 +1,39 @@ +--- +name: bisect-triton-release +description: Bisecting triton releases to figure out which commit causes regressions. +--- + +TokenSpeed depends on `tokenspeed-triton`, a vendor release of Triton, for +faster access to latest features. +If a new release causes a regression, we need to switch to upstream Triton and +bisect which commits caused the issue. + +## Overall Steps + +Always activate the local venv from the `tokenspeed/` directory to make sure +installing `triton` into it. + +1. Record the version of currently installed `triton` (or `triton-rocm` for AMD) + pip package, as a dependency to `torch`, so we can restore it later. +2. Patch `_triton.py` to switch `tokenspeed_triton` imports to direct `triton` + imports. +3. Ask which test in TokenSpeed shows regression. +4. Ask which upstream Triton commit is the last known good, and which upstream + Triton commit shows regression. +5. Build and install triton at the two commits gotten in previous step to + confirm that the regression happened due to upstream commit. + If not, then ask where the downstream triton repo is to inspect downstream + changes. +6. Once confirmed a upstream commit causes regression, start `git bisect` flow + to identify the exact root cause. +7. Once identified, reinstall the recorded `triton` (or `triton-rocm` package) + to the original version and revert local changes. + +## Build and Install Triton + +In Triton codebase, do the following + +```shell +pip install -r python/requirements.txt +pip install . +``` diff --git a/.skills/optimize-amd-gpu-kernel.md b/.skills/optimize-amd-gpu-kernel.md new file mode 100644 index 0000000..ab0e043 --- /dev/null +++ b/.skills/optimize-amd-gpu-kernel.md @@ -0,0 +1,91 @@ +--- +name: optimize-amd-gpu-kernel +description: Optimizing kernel performance on AMD Instinct MI GPUs. +--- + +## General Methodology + +* Start from representative model shapes. +* Understand whether the case is compute-, memory-, launch-, or latency-bound. + Then decide whether to optimize compute/memory overlap, memory movement, + kernel count, tune tiling and configurations, etc. +* Profile first, focus on the top bottleneck, change one thing, and rerun the + same benchmark/profiling method. +* Use Gluon for explicit low-level control: buffer load/store, async copy to + LDS, shared layouts, MFMA layout, wave count, and LLVM attributes. +* Pay attention to both `ttg` level and `llvm` level opportunities and balances. +* Tune configuration parameters, but do not overfit with many one-off switch + cases. + +## Profiling Tools + +* Use Proton for high-level TFLOp/s or TB/s calculation; check in code changes + for `triton.jit`/`gluon.jit` `repr` for reuse. +* Use `rocprofv3` in ROCm to understand low-level internals like counters. +* Proton also supports fine-grained profiling with `scope` APIs and + instrumentation mode. + +## Problem Approaches + +### General optimizations + +Applicable to various problems: + +* Prefer to launch enough workgroups to fill the GPU. +* Ensure proper software pipelining to break dependencies in the same loop + iteration. +* Prefer coalesced and vectorized async global memory load/store. +* If indexing range allows, prefer buffer load/store intrinsics in Gluon to + avoid out-of-bound branches and overheads. +* Avoid shared memory bank conflict if possible. Use padding instead of + swizzling. +* For async copy to LDS, arrange global load layouts so each thread issues wide, + aligned loads where possible; 128-bit per-thread loads are a good target. + +### Compute bound problems + +The key is to keep issuing MFMA instructions preferably every cycle, and avoid +exposed memory instruction latencies. Generally two approaches: + +* Use 4 waves per workgroup, and perform fine-grained per-instruction level + interleaving in the same wave on one SIMD. Typically needs controlling LLVM + knobs; can use `HIPOptions.llvm_fn_attrs`. +* Use 8 waves per workgroup, and perform coarse-grained multi-instruction level + interleaving across 2 waves on the same SIMD, to make sure those two waves + "ping-pong" among each other to overlap. Available via the + `amd.warp_pipeline_stage` API. + +Search and read AMD ISA docs and Triton codebase and examples to get +inspiration. + +* If high VGPR pressure, consider slice along M/N in the hot loop and interleave + to retire certain slices of loaded values earlier. + +### Memory bound problems + +The key is to saturate GPU memory bandwidth with enough inflight memory +instructions, and avoid exposed compute instruction cycles. + +* Prefetch using async load with higher number of shared memory buffers. +* Use double or triple buffering only when it hides real latency. Extra buffers + increase LDS/register pressure and may reduce occupancy or compiler quality. +* Use cache modifiers like `".cg"`, `".wt"`, etc. to control whether to cache at + certain levels. + +### Latency bound problems + +* Fuse multiple small kernels into one kernel when possible. + +### Small problem sizes + +* Perform split-k style optimization and launch second reduction kernel to + see if beneficial. +* Split-K can increase occupancy for high K, but the second reduction/finalize + kernel costs several microseconds. Only route it when it is consistently + faster than torch for the real shapes. + +### Expensive epilogue + +* Perform persistent kernel style optimization to see if beneficial. +* If an epilogue/reduction is unavoidable, consider whether it can be fused into + the producer without hurting occupancy or memory behavior. diff --git a/ACKNOWLEDGEMENTS.md b/ACKNOWLEDGEMENTS.md new file mode 100644 index 0000000..ff2b1be --- /dev/null +++ b/ACKNOWLEDGEMENTS.md @@ -0,0 +1,48 @@ +# Acknowledgements + +TokenSpeed builds on a broad open source ecosystem. We are grateful to the +projects and contributors whose work is acknowledged in our package-level +third-party notices. + +## tokenspeed + +The `tokenspeed` Python package acknowledges and thanks: + +- `NVIDIA/TensorRT-LLM` +- `meituan-longcat/SGLang-FluentLLM` +- `vllm-project/vllm` +- `ModelTC/lightllm` +- `sgl-project/sglang` + +## tokenspeed-kernel + +The `tokenspeed-kernel` package acknowledges and thanks: + +- `NVIDIA/TensorRT-LLM` +- `triton-lang/triton` +- `deepseek-ai/DeepEP` +- `deepseek-ai/DeepGEMM` +- `Dao-AILab/flash-attention` +- `deepseek-ai/FlashMLA` +- `Dao-AILab/fast-hadamard-transform` +- `sgl-kernel` +- `flashinfer-ai/flashinfer` + +## tokenspeed-mla + +The `tokenspeed-mla` package acknowledges and thanks: + +- `NVIDIA/cutlass` +- `flashinfer-ai/flashinfer` + +## License Notices + +We follow open source license requirements by keeping package-level +`THIRDPARTYNOTICES` files with the relevant third-party notices and license +texts. For Python packages built with setuptools, these files are included in +the package metadata with: + +```toml +[tool.setuptools] +license-files = ["LICENSE", "THIRDPARTYNOTICES"] +``` diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..7b13d5c --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,44 @@ +# General Agent Guidelines + +> If a `AGENTS.local.md` file exists alongside this file, read and respect it-- +> it contains developer-specific overrides that supplement this shared guidance. + +## Development environment + +* Before any work, check local Python venv and activate if one exists. +* Don't install pip packages outside the local Python venv if one exists. + +## Code changes + +* Add tests and update docs for the changed code. +* Before creating commits, run `pre-commit run --all-files` to format. +* Do not substitute a narrower lint command for the repository hook before + committing. Always run the exact `pre-commit run --all-files` command and + commit any formatter changes it makes. +* When creating commits, perform sign off on behalf of the author. + +## Dependency boundaries + +* `tokenspeed` runtime dependencies should stay vendor-neutral. +* Runtime code should use `tokenspeed-kernel` as its only kernel package + boundary. +* Third-party kernel libraries belong under `tokenspeed-kernel`; avoid direct + runtime dependencies or imports that bypass it. +* If a dependency repeatedly breaks during version upgrades or slows project + progress, consider removing it entirely or at least making it optional. + +## tokenspeed-kernel + +Inside the root tokenspeed-kernel/ directory: + +* All direct tokenspeed-triton imports should happen in `_triton.py` and then + re-import to other places. +* All direct third-party code should be placed in `thirdparty/` and imported + into `ops/` then registered via `register_kernel`. +* Prefer CuteDSL for NVIDIA GPU kernels and Triton Gluon for AMD GPU kernels. + Use Triton for portable solutions across vendors. Vendor libraries should + stay optional, and other solutions may be used as temporary transitions, but + new work should consolidate toward these backend choices. +* Files under `ops/` should follow `/` structure, like + `gemm/trtllm.py` or `attention/triton/`. +* When defining new public APIs, explain arguments and returns in docstring. diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..29984d4 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,4 @@ +# CLAUDE.md + +@AGENTS.md +@AGENTS.local.md diff --git a/CODE_OF_CONDUCT.md b/CODE_OF_CONDUCT.md new file mode 100644 index 0000000..ff8b20d --- /dev/null +++ b/CODE_OF_CONDUCT.md @@ -0,0 +1,128 @@ +# Contributor Covenant Code of Conduct + +## Our Pledge + +We as members, contributors, and leaders pledge to make participation in our +community a harassment-free experience for everyone, regardless of age, body +size, visible or invisible disability, ethnicity, sex characteristics, gender +identity and expression, level of experience, education, socio-economic status, +nationality, personal appearance, race, religion, or sexual identity +and orientation. + +We pledge to act and interact in ways that contribute to an open, welcoming, +diverse, inclusive, and healthy community. + +## Our Standards + +Examples of behavior that contributes to a positive environment for our +community include: + +* Demonstrating empathy and kindness toward other people +* Being respectful of differing opinions, viewpoints, and experiences +* Giving and gracefully accepting constructive feedback +* Accepting responsibility and apologizing to those affected by our mistakes, + and learning from the experience +* Focusing on what is best not just for us as individuals, but for the + overall community + +Examples of unacceptable behavior include: + +* The use of sexualized language or imagery, and sexual attention or + advances of any kind +* Trolling, insulting or derogatory comments, and personal or political attacks +* Public or private harassment +* Publishing others' private information, such as a physical or email + address, without their explicit permission +* Other conduct which could reasonably be considered inappropriate in a + professional setting + +## Enforcement Responsibilities + +Community leaders are responsible for clarifying and enforcing our standards of +acceptable behavior and will take appropriate and fair corrective action in +response to any behavior that they deem inappropriate, threatening, offensive, +or harmful. + +Community leaders have the right and responsibility to remove, edit, or reject +comments, commits, code, wiki edits, issues, and other contributions that are +not aligned to this Code of Conduct, and will communicate reasons for moderation +decisions when appropriate. + +## Scope + +This Code of Conduct applies within all community spaces, and also applies when +an individual is officially representing the community in public spaces. +Examples of representing our community include using an official e-mail address, +posting via an official social media account, or acting as an appointed +representative at an online or offline event. + +## Enforcement + +Instances of abusive, harassing, or otherwise unacceptable behavior may be +reported to the community leaders responsible for enforcement at +LightSeek Foundation . +All complaints will be reviewed and investigated promptly and fairly. + +All community leaders are obligated to respect the privacy and security of the +reporter of any incident. + +## Enforcement Guidelines + +Community leaders will follow these Community Impact Guidelines in determining +the consequences for any action they deem in violation of this Code of Conduct: + +### 1. Correction + +**Community Impact**: Use of inappropriate language or other behavior deemed +unprofessional or unwelcome in the community. + +**Consequence**: A private, written warning from community leaders, providing +clarity around the nature of the violation and an explanation of why the +behavior was inappropriate. A public apology may be requested. + +### 2. Warning + +**Community Impact**: A violation through a single incident or series +of actions. + +**Consequence**: A warning with consequences for continued behavior. No +interaction with the people involved, including unsolicited interaction with +those enforcing the Code of Conduct, for a specified period of time. This +includes avoiding interactions in community spaces as well as external channels +like social media. Violating these terms may lead to a temporary or +permanent ban. + +### 3. Temporary Ban + +**Community Impact**: A serious violation of community standards, including +sustained inappropriate behavior. + +**Consequence**: A temporary ban from any sort of interaction or public +communication with the community for a specified period of time. No public or +private interaction with the people involved, including unsolicited interaction +with those enforcing the Code of Conduct, is allowed during this period. +Violating these terms may lead to a permanent ban. + +### 4. Permanent Ban + +**Community Impact**: Demonstrating a pattern of violation of community +standards, including sustained inappropriate behavior, harassment of an +individual, or aggression toward or disparagement of classes of individuals. + +**Consequence**: A permanent ban from any sort of public interaction within +the community. + +## Attribution + +This Code of Conduct is adapted from the [Contributor Covenant][homepage], +version 2.0, available at +https://www.contributor-covenant.org/version/2/0/code_of_conduct.html. + +Community Impact Guidelines were inspired by [Mozilla's code of conduct +enforcement ladder](https://github.com/mozilla/diversity). + +[homepage]: https://www.contributor-covenant.org + +For answers to common questions about this code of conduct, see the FAQ at +https://www.contributor-covenant.org/faq. Translations are available at +https://www.contributor-covenant.org/translations. diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..3c727b4 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,38 @@ +# TokenSpeed Project Ethos and External Contributions + +TokenSpeed is intentionally built as a lean, engineering-focused infrastructure project. + +In spirit, our development style is closer to projects like OpenAI Triton: +a lean and small core team, move fast, high technical ownership, and an extremely high quality bar. + +As a result, we are intentionally selective about external contributions — not because we do not value community involvement, but because maintaining long-term simplicity, performance, and reliability matters deeply for infrastructure software. + +TokenSpeed is built by contributors from different companies and open-source communities under the LightSeek Foundation ecosystem. Our mission is simple: + +We believe the future of AI should be transparent, collaborative, and inclusive. LightSeek exists to accelerate the development of open, efficient, and trustworthy AI systems—so that progress in AI benefits everyone, not just a few. + +## We Welcome External Contributions + +We especially welcome: + +- obvious bug fixes +- small and verifiable production-needed features +- performance optimizations that fit the existing codebase style and do not introduce unnecessary complexity +- documentation, tooling, and benchmarking improvements + +For larger features or architectural changes, we generally recommend starting with an RFC or design discussion first before implementation. + +## Engineering Principles + +We value: + +- reproducible performance +- correctness and stability +- production-oriented engineering +- maintainable implementations + +We believe open-source infrastructure advances through open collaboration, healthy technical discussion, and shared engineering progress. + +Thanks to everyone contributing code, ideas, reviews, benchmarks, and feedback. + +— TokenSpeed Team diff --git a/GOVERNANCE.md b/GOVERNANCE.md new file mode 100644 index 0000000..62e043c --- /dev/null +++ b/GOVERNANCE.md @@ -0,0 +1,42 @@ +# TokenSpeed Governance + +TokenSpeed is maintained by the LightSeek Foundation and the project core +maintainers. + +## Core Maintainers + +Named core maintainers include: + +- @zhyncs +- @borontion +- @antiagainst +- @dongjiyingdjy +- @syuoni +- @LorrinWWW + +This list is not exhaustive and may change as maintainers are added or step +back. Core maintainers are responsible for project direction, code review, +releases, and long-term maintenance. + +## Decisions + +Most decisions are made through GitHub issues, pull requests, and review +discussion. Maintainers should seek consensus when possible. + +Major project decisions, including maintainer membership, release policy, +project scope, and changes to this document, require approval from more than +half of the active core maintainers. + +## Adding Maintainers + +A contributor may be nominated as a core maintainer after making significant +contributions to TokenSpeed and helping maintain the project for at least three +months. + +A nomination must be made by an existing core maintainer and passes when more +than half of the active core maintainers approve it. + +## Governance Changes + +Changes to this document should be proposed by pull request and require +approval from more than half of the active core maintainers. diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..e95f89c --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 LightSeek Foundation + +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. diff --git a/README.md b/README.md new file mode 100644 index 0000000..2ed3362 --- /dev/null +++ b/README.md @@ -0,0 +1,45 @@ +

+ TokenSpeed: Tokens at the speed of light +

+ +TokenSpeed is a speed-of-light LLM inference engine designed for **agentic workloads**, with TensorRT-LLM-level performance and vLLM-level usability. Our goal is to be the most performant inference engine for production agentic workloads. + +Core components: + +- **Modeling layer**: local-SPMD design with a static compiler that generates + collective communication from module-boundary placement annotations, so users + do not hand-write parallelism logic. +- **Scheduler**: C++ control plane and Python execution plane. Request + lifecycle, KV cache ownership, and overlap timing are encoded as a + finite-state machine, with safe KV resource reuse enforced by the type system at compile time. +- **Kernels**: pluggable, layered kernel system with a portable public API and + a centralized registry including one of the fastest **MLA** + (Multi-head Latent Attention) implementations on Blackwell for agentic workload. +- **Entrypoint**: SMG-integrated AsyncLLM for low-overhead CPU-side request + handling. + +## News + +- [2026/06] Deep dive into the design and optimization of TokenSpeed-Kernel. [[blog](https://pytorch.org/blog/lightseek-tokenspeed-kernel/)] +- [2026/05] 🚀 TokenSpeed hits 580 TPS on Qwen3.5-397B-A17B for agentic workloads. [[blog](https://pytorch.org/blog/up-to-580tps-new-speed-record-of-qwen3-5-397b-a17b-on-gpu-for-agentic-workloads-with-tokenspeed/)] +- [2026/05] TokenSpeed announced — a speed-of-light LLM inference engine for agentic workloads. [[blog](https://lightseek.org/blog/lightseek-tokenspeed.html)] + +## Blogs and Talks + +For technical blogs, conference talks, and engineering articles from LightSeek Foundation, visit the [LightSeek Blog](https://lightseek.org/blog/). + +## Performance Comparison + +TokenSpeed vs. TensorRT-LLM Pareto curves on agentic workload (Kimi K2.5, B200) + +## Documentation + +Start here: + +- [Docs Index](https://lightseek.org/tokenspeed/) +- [Getting Started](https://lightseek.org/tokenspeed/guides/getting-started) +- [Launching a Server](https://lightseek.org/tokenspeed/guides/launching) +- [Model Recipes](https://lightseek.org/tokenspeed/recipes/models) +- [Server Parameters](https://lightseek.org/tokenspeed/configuration/server) +- [Compatible Parameters](https://lightseek.org/tokenspeed/configuration/compatible-parameters) +- [Parallelism](https://lightseek.org/tokenspeed/serving/parallelism) diff --git a/README.wehub.md b/README.wehub.md new file mode 100644 index 0000000..d9e5479 --- /dev/null +++ b/README.wehub.md @@ -0,0 +1,7 @@ +# WeHub 来源说明 + +- 原始项目：`lightseekorg/tokenspeed` +- 原始仓库：https://github.com/lightseekorg/tokenspeed +- 导入方式：上游默认分支的最新快照 +- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准 +- 本文件仅用于记录来源，不代表 WeHub 是原项目作者 diff --git a/SECURITY.md b/SECURITY.md new file mode 100644 index 0000000..f6cb141 --- /dev/null +++ b/SECURITY.md @@ -0,0 +1,7 @@ +# Security Policy + +## Reporting a Vulnerability + +If you discover a security vulnerability in TokenSpeed, please report it through GitHub Security Advisories or contact the maintainers directly at [contact@lightseek.org](mailto:contact@lightseek.org). + +We will acknowledge all valid reports promptly, investigate the issue, and work to provide a fix as quickly as possible. We appreciate your responsible disclosure and efforts to help keep TokenSpeed secure. diff --git a/assets/banner/tokenspeed-banner.png b/assets/banner/tokenspeed-banner.png new file mode 100644 index 0000000..1b63faa Binary files /dev/null and b/assets/banner/tokenspeed-banner.png differ diff --git a/assets/perf/tokenspeed-kimi-k2.5-performance.png b/assets/perf/tokenspeed-kimi-k2.5-performance.png new file mode 100644 index 0000000..459fa2c Binary files /dev/null and b/assets/perf/tokenspeed-kimi-k2.5-performance.png differ diff --git a/docker/Dockerfile b/docker/Dockerfile new file mode 100644 index 0000000..016781a --- /dev/null +++ b/docker/Dockerfile @@ -0,0 +1,24 @@ +ARG RUNNER_TAG=cu130-torch-2.11.0 +FROM lightseekorg/tokenspeed-runner:${RUNNER_TAG} + +USER root + +ARG MAX_JOBS=16 +ENV PIP_BREAK_SYSTEM_PACKAGES=1 + +WORKDIR /workspace + +COPY . /workspace + +# libnuma1: runtime dep of torch_memory_saver (the Sleep/Wake Up memory saver), +# which is a hard dependency of the python package; the base runner image lacks it. +RUN apt-get update && apt-get install -y --no-install-recommends libssl-dev libopenmpi-dev libnuma1 && \ + rm -rf /var/lib/apt/lists/* + +ARG CUDA_ARCH_LIST="9.0a 10.0a" +RUN MAX_JOBS=${MAX_JOBS} FLASHINFER_CUDA_ARCH_LIST="${CUDA_ARCH_LIST}" TOKENSPEED_KERNEL_BACKEND=cuda \ + pip install tokenspeed-kernel/python/ --no-build-isolation && \ + pip install tokenspeed-scheduler/ && \ + pip install "./python" --no-build-isolation + +CMD ["/bin/bash"] diff --git a/docs/.gitignore b/docs/.gitignore new file mode 100644 index 0000000..c76bc00 --- /dev/null +++ b/docs/.gitignore @@ -0,0 +1,3 @@ +node_modules/ +.vitepress/cache/ +.vitepress/dist/ diff --git a/docs/.vitepress/config.mts b/docs/.vitepress/config.mts new file mode 100644 index 0000000..1289eb3 --- /dev/null +++ b/docs/.vitepress/config.mts @@ -0,0 +1,84 @@ +import { defineConfig } from "vitepress"; + +export default defineConfig({ + lang: "en-US", + title: "TokenSpeed", + description: "TokenSpeed is a speed-of-light LLM inference engine.", + base: "/tokenspeed/", + srcDir: ".", + cleanUrls: true, + lastUpdated: true, + head: [ + ["meta", { property: "og:title", content: "TokenSpeed Docs" }], + [ + "meta", + { + property: "og:description", + content: + "Guides and reference for running TokenSpeed, configuring the server, and operating multi-GPU inference." + } + ] + ], + themeConfig: { + nav: [ + { text: "Docs Home", link: "/" }, + { text: "Getting Started", link: "/guides/getting-started" }, + { text: "Launching", link: "/guides/launching" }, + { text: "Recipes", link: "/recipes/models" }, + { text: "Configuration", link: "/configuration/server" }, + { text: "GitHub", link: "https://github.com/lightseekorg/tokenspeed" } + ], + sidebar: [ + { + text: "Overview", + items: [{ text: "Docs Home", link: "/" }] + }, + { + text: "Guides", + items: [ + { text: "Getting Started", link: "/guides/getting-started" }, + { text: "Launching a Server", link: "/guides/launching" } + ] + }, + { + text: "Configuration", + items: [ + { text: "Server Parameters", link: "/configuration/server" }, + { + text: "Compatible Parameters", + link: "/configuration/compatible-parameters" + } + ] + }, + { + text: "Recipes", + items: [{ text: "Model Recipes", link: "/recipes/models" }] + }, + { + text: "Serving", + items: [{ text: "Parallelism", link: "/serving/parallelism" }] + } + ], + search: { + provider: "local" + }, + editLink: { + pattern: + "https://github.com/lightseekorg/tokenspeed/edit/main/docs/:path", + text: "Edit this page on GitHub" + }, + footer: { + message: "TokenSpeed documentation", + copyright: "Copyright © lightseekorg" + }, + socialLinks: [ + { icon: "github", link: "https://github.com/lightseekorg/tokenspeed" } + ], + outline: "deep", + outlineTitle: "On this page", + docFooter: { + prev: "Previous page", + next: "Next page" + } + } +}); diff --git a/docs/bun.lock b/docs/bun.lock new file mode 100644 index 0000000..bef3332 --- /dev/null +++ b/docs/bun.lock @@ -0,0 +1,359 @@ +{ + "lockfileVersion": 1, + "configVersion": 1, + "workspaces": { + "": { + "name": "tokenspeed-docs", + "devDependencies": { + "vitepress": "^1.6.3", + }, + }, + }, + "packages": { + "@algolia/abtesting": ["@algolia/abtesting@1.16.2", "", { "dependencies": { "@algolia/client-common": "5.50.2", "@algolia/requester-browser-xhr": "5.50.2", "@algolia/requester-fetch": "5.50.2", "@algolia/requester-node-http": "5.50.2" } }, "sha512-n9s6bEV6imdtIEd+BGP7WkA4pEZ5YTdgQ05JQhHwWawHg3hyjpNwC0TShGz6zWhv+jfLDGA/6FFNbySFS0P9cw=="], + + "@algolia/autocomplete-core": ["@algolia/autocomplete-core@1.17.7", "", { "dependencies": { "@algolia/autocomplete-plugin-algolia-insights": "1.17.7", "@algolia/autocomplete-shared": "1.17.7" } }, "sha512-BjiPOW6ks90UKl7TwMv7oNQMnzU+t/wk9mgIDi6b1tXpUek7MW0lbNOUHpvam9pe3lVCf4xPFT+lK7s+e+fs7Q=="], + + "@algolia/autocomplete-plugin-algolia-insights": ["@algolia/autocomplete-plugin-algolia-insights@1.17.7", "", { "dependencies": { "@algolia/autocomplete-shared": "1.17.7" }, "peerDependencies": { "search-insights": ">= 1 < 3" } }, "sha512-Jca5Ude6yUOuyzjnz57og7Et3aXjbwCSDf/8onLHSQgw1qW3ALl9mrMWaXb5FmPVkV3EtkD2F/+NkT6VHyPu9A=="], + + "@algolia/autocomplete-preset-algolia": ["@algolia/autocomplete-preset-algolia@1.17.7", "", { "dependencies": { "@algolia/autocomplete-shared": "1.17.7" }, "peerDependencies": { "@algolia/client-search": ">= 4.9.1 < 6", "algoliasearch": ">= 4.9.1 < 6" } }, "sha512-ggOQ950+nwbWROq2MOCIL71RE0DdQZsceqrg32UqnhDz8FlO9rL8ONHNsI2R1MH0tkgVIDKI/D0sMiUchsFdWA=="], + + "@algolia/autocomplete-shared": ["@algolia/autocomplete-shared@1.17.7", "", { "peerDependencies": { "@algolia/client-search": ">= 4.9.1 < 6", "algoliasearch": ">= 4.9.1 < 6" } }, "sha512-o/1Vurr42U/qskRSuhBH+VKxMvkkUVTLU6WZQr+L5lGZZLYWyhdzWjW0iGXY7EkwRTjBqvN2EsR81yCTGV/kmg=="], + + "@algolia/client-abtesting": ["@algolia/client-abtesting@5.50.2", "", { "dependencies": { "@algolia/client-common": "5.50.2", "@algolia/requester-browser-xhr": "5.50.2", "@algolia/requester-fetch": "5.50.2", "@algolia/requester-node-http": "5.50.2" } }, "sha512-52iq0vHy1sphgnwoZyx5PmbEt8hsh+m7jD123LmBs6qy4GK7LbYZIeKd+nSnSipN2zvKRZ2zScS6h9PW3J7SXg=="], + + "@algolia/client-analytics": ["@algolia/client-analytics@5.50.2", "", { "dependencies": { "@algolia/client-common": "5.50.2", "@algolia/requester-browser-xhr": "5.50.2", "@algolia/requester-fetch": "5.50.2", "@algolia/requester-node-http": "5.50.2" } }, "sha512-WpPIUg+cSG2aPUG0gS8Ko9DwRgbRPUZxJkolhL2aCsmSlcEEZT65dILrfg5ovcxtx0Kvr+xtBVsTMtsQWRtPDQ=="], + + "@algolia/client-common": ["@algolia/client-common@5.50.2", "", {}, "sha512-Gj2MgtArGcsr82kIqRlo6/dCAFjrs2gLByEqyRENuT7ugrSMFuqg1vDzeBjRL1t3EJEJCFtT0PLX3gB8A6Hq4Q=="], + + "@algolia/client-insights": ["@algolia/client-insights@5.50.2", "", { "dependencies": { "@algolia/client-common": "5.50.2", "@algolia/requester-browser-xhr": "5.50.2", "@algolia/requester-fetch": "5.50.2", "@algolia/requester-node-http": "5.50.2" } }, "sha512-CUqoid5jDpmrc0oK3/xuZXFt6kwT0P9Lw7/nsM14YTr6puvmi+OUKmURpmebQF22S2vCG8L1DAoXXujxQUi/ug=="], + + "@algolia/client-personalization": ["@algolia/client-personalization@5.50.2", "", { "dependencies": { "@algolia/client-common": "5.50.2", "@algolia/requester-browser-xhr": "5.50.2", "@algolia/requester-fetch": "5.50.2", "@algolia/requester-node-http": "5.50.2" } }, "sha512-AndZWFoc0gbP5901OeQJ73BazgGgSGiBEba4ohdoJuZwHTO2Gio8Q4L1VLmytMBYcviVigB0iICToMvEJxI4ug=="], + + "@algolia/client-query-suggestions": ["@algolia/client-query-suggestions@5.50.2", "", { "dependencies": { "@algolia/client-common": "5.50.2", "@algolia/requester-browser-xhr": "5.50.2", "@algolia/requester-fetch": "5.50.2", "@algolia/requester-node-http": "5.50.2" } }, "sha512-NWoL+psEkz5dIzweaByVXuEB45wS8/rk0E0AhMMnaVJdVs7TcACPH2/OURm+N0xRDITkTHqCna823rd6Uqntdg=="], + + "@algolia/client-search": ["@algolia/client-search@5.50.2", "", { "dependencies": { "@algolia/client-common": "5.50.2", "@algolia/requester-browser-xhr": "5.50.2", "@algolia/requester-fetch": "5.50.2", "@algolia/requester-node-http": "5.50.2" } }, "sha512-ypSboUJ3XJoQz5DeDo82hCnrRuwq3q9ZdFhVKAik9TnZh1DvLqoQsrbBjXg7C7zQOtV/Qbge/HmyoV6V5L7MhQ=="], + + "@algolia/ingestion": ["@algolia/ingestion@1.50.2", "", { "dependencies": { "@algolia/client-common": "5.50.2", "@algolia/requester-browser-xhr": "5.50.2", "@algolia/requester-fetch": "5.50.2", "@algolia/requester-node-http": "5.50.2" } }, "sha512-VlR2FRXLw2bCB94SQo6zxg/Qi+547aOji6Pb+dKE7h1DMCCY317St+OpjpmgzE+bT2O9ALIc0V4nVIBOd7Gy+Q=="], + + "@algolia/monitoring": ["@algolia/monitoring@1.50.2", "", { "dependencies": { "@algolia/client-common": "5.50.2", "@algolia/requester-browser-xhr": "5.50.2", "@algolia/requester-fetch": "5.50.2", "@algolia/requester-node-http": "5.50.2" } }, "sha512-Cmvfp2+qopzQt8OilU97rhLhosq7ZrB6uieok3EwFUqG/aalPg6DgfCmu0yJMrYe+KMC1qRVt1MTRAUwLknUMQ=="], + + "@algolia/recommend": ["@algolia/recommend@5.50.2", "", { "dependencies": { "@algolia/client-common": "5.50.2", "@algolia/requester-browser-xhr": "5.50.2", "@algolia/requester-fetch": "5.50.2", "@algolia/requester-node-http": "5.50.2" } }, "sha512-jrkuyKoOM7dFWQ/6Y4hQAse2SC3L/RldG6GnPjMvAj65h+7Ubb51S0pKk4ofSStF0xm4LCNe0C4T6XX4nOFDiQ=="], + + "@algolia/requester-browser-xhr": ["@algolia/requester-browser-xhr@5.50.2", "", { "dependencies": { "@algolia/client-common": "5.50.2" } }, "sha512-4107YLJqCudPiBUlwnk6oTSUVwU7ab+qL1SfQGEDYI8DZH5gsf1ekPt9JykXRKYXf2IfouFL5GiCY/PHTFIjYw=="], + + "@algolia/requester-fetch": ["@algolia/requester-fetch@5.50.2", "", { "dependencies": { "@algolia/client-common": "5.50.2" } }, "sha512-vOrd3MQpLgmf6wXAueTuZ/cA0W4uRwIHHaxNy3h+a6YcNn6bCV/gFdZuv3F13v593zRU2k5R75NmvRWLenvMrw=="], + + "@algolia/requester-node-http": ["@algolia/requester-node-http@5.50.2", "", { "dependencies": { "@algolia/client-common": "5.50.2" } }, "sha512-Mu9BFtgzGqDUy5Bcs2nMyoILIFSN13GKQaklKAFIsd0K3/9CpNyfeBc+/+Qs6mFZLlxG9qzullO7h+bjcTBuGQ=="], + + "@babel/helper-string-parser": ["@babel/helper-string-parser@7.27.1", "", {}, "sha512-qMlSxKbpRlAridDExk92nSobyDdpPijUq2DW6oDnUqd0iOGxmQjyqhMIihI9+zv4LPyZdRje2cavWPbCbWm3eA=="], + + "@babel/helper-validator-identifier": ["@babel/helper-validator-identifier@7.28.5", "", {}, "sha512-qSs4ifwzKJSV39ucNjsvc6WVHs6b7S03sOh2OcHF9UHfVPqWWALUsNUVzhSBiItjRZoLHx7nIarVjqKVusUZ1Q=="], + + "@babel/parser": ["@babel/parser@7.29.2", "", { "dependencies": { "@babel/types": "^7.29.0" }, "bin": "./bin/babel-parser.js" }, "sha512-4GgRzy/+fsBa72/RZVJmGKPmZu9Byn8o4MoLpmNe1m8ZfYnz5emHLQz3U4gLud6Zwl0RZIcgiLD7Uq7ySFuDLA=="], + + "@babel/types": ["@babel/types@7.29.0", "", { "dependencies": { "@babel/helper-string-parser": "^7.27.1", "@babel/helper-validator-identifier": "^7.28.5" } }, "sha512-LwdZHpScM4Qz8Xw2iKSzS+cfglZzJGvofQICy7W7v4caru4EaAmyUuO6BGrbyQ2mYV11W0U8j5mBhd14dd3B0A=="], + + "@docsearch/css": ["@docsearch/css@3.8.2", "", {}, "sha512-y05ayQFyUmCXze79+56v/4HpycYF3uFqB78pLPrSV5ZKAlDuIAAJNhaRi8tTdRNXh05yxX/TyNnzD6LwSM89vQ=="], + + "@docsearch/js": ["@docsearch/js@3.8.2", "", { "dependencies": { "@docsearch/react": "3.8.2", "preact": "^10.0.0" } }, "sha512-Q5wY66qHn0SwA7Taa0aDbHiJvaFJLOJyHmooQ7y8hlwwQLQ/5WwCcoX0g7ii04Qi2DJlHsd0XXzJ8Ypw9+9YmQ=="], + + "@docsearch/react": ["@docsearch/react@3.8.2", "", { "dependencies": { "@algolia/autocomplete-core": "1.17.7", "@algolia/autocomplete-preset-algolia": "1.17.7", "@docsearch/css": "3.8.2", "algoliasearch": "^5.14.2" }, "peerDependencies": { "@types/react": ">= 16.8.0 < 19.0.0", "react": ">= 16.8.0 < 19.0.0", "react-dom": ">= 16.8.0 < 19.0.0", "search-insights": ">= 1 < 3" }, "optionalPeers": ["@types/react", "react", "react-dom", "search-insights"] }, "sha512-xCRrJQlTt8N9GU0DG4ptwHRkfnSnD/YpdeaXe02iKfqs97TkZJv60yE+1eq/tjPcVnTW8dP5qLP7itifFVV5eg=="], + + "@esbuild/aix-ppc64": ["@esbuild/aix-ppc64@0.21.5", "", { "os": "aix", "cpu": "ppc64" }, "sha512-1SDgH6ZSPTlggy1yI6+Dbkiz8xzpHJEVAlF/AM1tHPLsf5STom9rwtjE4hKAF20FfXXNTFqEYXyJNWh1GiZedQ=="], + + "@esbuild/android-arm": ["@esbuild/android-arm@0.21.5", "", { "os": "android", "cpu": "arm" }, "sha512-vCPvzSjpPHEi1siZdlvAlsPxXl7WbOVUBBAowWug4rJHb68Ox8KualB+1ocNvT5fjv6wpkX6o/iEpbDrf68zcg=="], + + "@esbuild/android-arm64": ["@esbuild/android-arm64@0.21.5", "", { "os": "android", "cpu": "arm64" }, "sha512-c0uX9VAUBQ7dTDCjq+wdyGLowMdtR/GoC2U5IYk/7D1H1JYC0qseD7+11iMP2mRLN9RcCMRcjC4YMclCzGwS/A=="], + + "@esbuild/android-x64": ["@esbuild/android-x64@0.21.5", "", { "os": "android", "cpu": "x64" }, "sha512-D7aPRUUNHRBwHxzxRvp856rjUHRFW1SdQATKXH2hqA0kAZb1hKmi02OpYRacl0TxIGz/ZmXWlbZgjwWYaCakTA=="], + + "@esbuild/darwin-arm64": ["@esbuild/darwin-arm64@0.21.5", "", { "os": "darwin", "cpu": "arm64" }, "sha512-DwqXqZyuk5AiWWf3UfLiRDJ5EDd49zg6O9wclZ7kUMv2WRFr4HKjXp/5t8JZ11QbQfUS6/cRCKGwYhtNAY88kQ=="], + + "@esbuild/darwin-x64": ["@esbuild/darwin-x64@0.21.5", "", { "os": "darwin", "cpu": "x64" }, "sha512-se/JjF8NlmKVG4kNIuyWMV/22ZaerB+qaSi5MdrXtd6R08kvs2qCN4C09miupktDitvh8jRFflwGFBQcxZRjbw=="], + + "@esbuild/freebsd-arm64": ["@esbuild/freebsd-arm64@0.21.5", "", { "os": "freebsd", "cpu": "arm64" }, "sha512-5JcRxxRDUJLX8JXp/wcBCy3pENnCgBR9bN6JsY4OmhfUtIHe3ZW0mawA7+RDAcMLrMIZaf03NlQiX9DGyB8h4g=="], + + "@esbuild/freebsd-x64": ["@esbuild/freebsd-x64@0.21.5", "", { "os": "freebsd", "cpu": "x64" }, "sha512-J95kNBj1zkbMXtHVH29bBriQygMXqoVQOQYA+ISs0/2l3T9/kj42ow2mpqerRBxDJnmkUDCaQT/dfNXWX/ZZCQ=="], + + "@esbuild/linux-arm": ["@esbuild/linux-arm@0.21.5", "", { "os": "linux", "cpu": "arm" }, "sha512-bPb5AHZtbeNGjCKVZ9UGqGwo8EUu4cLq68E95A53KlxAPRmUyYv2D6F0uUI65XisGOL1hBP5mTronbgo+0bFcA=="], + + "@esbuild/linux-arm64": ["@esbuild/linux-arm64@0.21.5", "", { "os": "linux", "cpu": "arm64" }, "sha512-ibKvmyYzKsBeX8d8I7MH/TMfWDXBF3db4qM6sy+7re0YXya+K1cem3on9XgdT2EQGMu4hQyZhan7TeQ8XkGp4Q=="], + + "@esbuild/linux-ia32": ["@esbuild/linux-ia32@0.21.5", "", { "os": "linux", "cpu": "ia32" }, "sha512-YvjXDqLRqPDl2dvRODYmmhz4rPeVKYvppfGYKSNGdyZkA01046pLWyRKKI3ax8fbJoK5QbxblURkwK/MWY18Tg=="], + + "@esbuild/linux-loong64": ["@esbuild/linux-loong64@0.21.5", "", { "os": "linux", "cpu": "none" }, "sha512-uHf1BmMG8qEvzdrzAqg2SIG/02+4/DHB6a9Kbya0XDvwDEKCoC8ZRWI5JJvNdUjtciBGFQ5PuBlpEOXQj+JQSg=="], + + "@esbuild/linux-mips64el": ["@esbuild/linux-mips64el@0.21.5", "", { "os": "linux", "cpu": "none" }, "sha512-IajOmO+KJK23bj52dFSNCMsz1QP1DqM6cwLUv3W1QwyxkyIWecfafnI555fvSGqEKwjMXVLokcV5ygHW5b3Jbg=="], + + "@esbuild/linux-ppc64": ["@esbuild/linux-ppc64@0.21.5", "", { "os": "linux", "cpu": "ppc64" }, "sha512-1hHV/Z4OEfMwpLO8rp7CvlhBDnjsC3CttJXIhBi+5Aj5r+MBvy4egg7wCbe//hSsT+RvDAG7s81tAvpL2XAE4w=="], + + "@esbuild/linux-riscv64": ["@esbuild/linux-riscv64@0.21.5", "", { "os": "linux", "cpu": "none" }, "sha512-2HdXDMd9GMgTGrPWnJzP2ALSokE/0O5HhTUvWIbD3YdjME8JwvSCnNGBnTThKGEB91OZhzrJ4qIIxk/SBmyDDA=="], + + "@esbuild/linux-s390x": ["@esbuild/linux-s390x@0.21.5", "", { "os": "linux", "cpu": "s390x" }, "sha512-zus5sxzqBJD3eXxwvjN1yQkRepANgxE9lgOW2qLnmr8ikMTphkjgXu1HR01K4FJg8h1kEEDAqDcZQtbrRnB41A=="], + + "@esbuild/linux-x64": ["@esbuild/linux-x64@0.21.5", "", { "os": "linux", "cpu": "x64" }, "sha512-1rYdTpyv03iycF1+BhzrzQJCdOuAOtaqHTWJZCWvijKD2N5Xu0TtVC8/+1faWqcP9iBCWOmjmhoH94dH82BxPQ=="], + + "@esbuild/netbsd-x64": ["@esbuild/netbsd-x64@0.21.5", "", { "os": "none", "cpu": "x64" }, "sha512-Woi2MXzXjMULccIwMnLciyZH4nCIMpWQAs049KEeMvOcNADVxo0UBIQPfSmxB3CWKedngg7sWZdLvLczpe0tLg=="], + + "@esbuild/openbsd-x64": ["@esbuild/openbsd-x64@0.21.5", "", { "os": "openbsd", "cpu": "x64" }, "sha512-HLNNw99xsvx12lFBUwoT8EVCsSvRNDVxNpjZ7bPn947b8gJPzeHWyNVhFsaerc0n3TsbOINvRP2byTZ5LKezow=="], + + "@esbuild/sunos-x64": ["@esbuild/sunos-x64@0.21.5", "", { "os": "sunos", "cpu": "x64" }, "sha512-6+gjmFpfy0BHU5Tpptkuh8+uw3mnrvgs+dSPQXQOv3ekbordwnzTVEb4qnIvQcYXq6gzkyTnoZ9dZG+D4garKg=="], + + "@esbuild/win32-arm64": ["@esbuild/win32-arm64@0.21.5", "", { "os": "win32", "cpu": "arm64" }, "sha512-Z0gOTd75VvXqyq7nsl93zwahcTROgqvuAcYDUr+vOv8uHhNSKROyU961kgtCD1e95IqPKSQKH7tBTslnS3tA8A=="], + + "@esbuild/win32-ia32": ["@esbuild/win32-ia32@0.21.5", "", { "os": "win32", "cpu": "ia32" }, "sha512-SWXFF1CL2RVNMaVs+BBClwtfZSvDgtL//G/smwAc5oVK/UPu2Gu9tIaRgFmYFFKrmg3SyAjSrElf0TiJ1v8fYA=="], + + "@esbuild/win32-x64": ["@esbuild/win32-x64@0.21.5", "", { "os": "win32", "cpu": "x64" }, "sha512-tQd/1efJuzPC6rCFwEvLtci/xNFcTZknmXs98FYDfGE4wP9ClFV98nyKrzJKVPMhdDnjzLhdUyMX4PsQAPjwIw=="], + + "@iconify-json/simple-icons": ["@iconify-json/simple-icons@1.2.79", "", { "dependencies": { "@iconify/types": "*" } }, "sha512-aNyO7Fd1qej9oQfIyohYFRv0lhQLaZ+6UkK1c1qwax0MDPUOZOdq65MlU500kow97pD/W+b2u1And3e25eE24Q=="], + + "@iconify/types": ["@iconify/types@2.0.0", "", {}, "sha512-+wluvCrRhXrhyOmRDJ3q8mux9JkKy5SJ/v8ol2tu4FVjyYvtEzkc/3pK15ET6RKg4b4w4BmTk1+gsCUhf21Ykg=="], + + "@jridgewell/sourcemap-codec": ["@jridgewell/sourcemap-codec@1.5.5", "", {}, "sha512-cYQ9310grqxueWbl+WuIUIaiUaDcj7WOq5fVhEljNVgRfOUhY9fy2zTvfoqWsnebh8Sl70VScFbICvJnLKB0Og=="], + + "@rollup/rollup-android-arm-eabi": ["@rollup/rollup-android-arm-eabi@4.60.2", "", { "os": "android", "cpu": "arm" }, "sha512-dnlp69efPPg6Uaw2dVqzWRfAWRnYVb1XJ8CyyhIbZeaq4CA5/mLeZ1IEt9QqQxmbdvagjLIm2ZL8BxXv5lH4Yw=="], + + "@rollup/rollup-android-arm64": ["@rollup/rollup-android-arm64@4.60.2", "", { "os": "android", "cpu": "arm64" }, "sha512-OqZTwDRDchGRHHm/hwLOL7uVPB9aUvI0am/eQuWMNyFHf5PSEQmyEeYYheA0EPPKUO/l0uigCp+iaTjoLjVoHg=="], + + "@rollup/rollup-darwin-arm64": ["@rollup/rollup-darwin-arm64@4.60.2", "", { "os": "darwin", "cpu": "arm64" }, "sha512-UwRE7CGpvSVEQS8gUMBe1uADWjNnVgP3Iusyda1nSRwNDCsRjnGc7w6El6WLQsXmZTbLZx9cecegumcitNfpmA=="], + + "@rollup/rollup-darwin-x64": ["@rollup/rollup-darwin-x64@4.60.2", "", { "os": "darwin", "cpu": "x64" }, "sha512-gjEtURKLCC5VXm1I+2i1u9OhxFsKAQJKTVB8WvDAHF+oZlq0GTVFOlTlO1q3AlCTE/DF32c16ESvfgqR7343/g=="], + + "@rollup/rollup-freebsd-arm64": ["@rollup/rollup-freebsd-arm64@4.60.2", "", { "os": "freebsd", "cpu": "arm64" }, "sha512-Bcl6CYDeAgE70cqZaMojOi/eK63h5Me97ZqAQoh77VPjMysA/4ORQBRGo3rRy45x4MzVlU9uZxs8Uwy7ZaKnBw=="], + + "@rollup/rollup-freebsd-x64": ["@rollup/rollup-freebsd-x64@4.60.2", "", { "os": "freebsd", "cpu": "x64" }, "sha512-LU+TPda3mAE2QB0/Hp5VyeKJivpC6+tlOXd1VMoXV/YFMvk/MNk5iXeBfB4MQGRWyOYVJ01625vjkr0Az98OJQ=="], + + "@rollup/rollup-linux-arm-gnueabihf": ["@rollup/rollup-linux-arm-gnueabihf@4.60.2", "", { "os": "linux", "cpu": "arm" }, "sha512-2QxQrM+KQ7DAW4o22j+XZ6RKdxjLD7BOWTP0Bv0tmjdyhXSsr2Ul1oJDQqh9Zf5qOwTuTc7Ek83mOFaKnodPjg=="], + + "@rollup/rollup-linux-arm-musleabihf": ["@rollup/rollup-linux-arm-musleabihf@4.60.2", "", { "os": "linux", "cpu": "arm" }, "sha512-TbziEu2DVsTEOPif2mKWkMeDMLoYjx95oESa9fkQQK7r/Orta0gnkcDpzwufEcAO2BLBsD7mZkXGFqEdMRRwfw=="], + + "@rollup/rollup-linux-arm64-gnu": ["@rollup/rollup-linux-arm64-gnu@4.60.2", "", { "os": "linux", "cpu": "arm64" }, "sha512-bO/rVDiDUuM2YfuCUwZ1t1cP+/yqjqz+Xf2VtkdppefuOFS2OSeAfgafaHNkFn0t02hEyXngZkxtGqXcXwO8Rg=="], + + "@rollup/rollup-linux-arm64-musl": ["@rollup/rollup-linux-arm64-musl@4.60.2", "", { "os": "linux", "cpu": "arm64" }, "sha512-hr26p7e93Rl0Za+JwW7EAnwAvKkehh12BU1Llm9Ykiibg4uIr2rbpxG9WCf56GuvidlTG9KiiQT/TXT1yAWxTA=="], + + "@rollup/rollup-linux-loong64-gnu": ["@rollup/rollup-linux-loong64-gnu@4.60.2", "", { "os": "linux", "cpu": "none" }, "sha512-pOjB/uSIyDt+ow3k/RcLvUAOGpysT2phDn7TTUB3n75SlIgZzM6NKAqlErPhoFU+npgY3/n+2HYIQVbF70P9/A=="], + + "@rollup/rollup-linux-loong64-musl": ["@rollup/rollup-linux-loong64-musl@4.60.2", "", { "os": "linux", "cpu": "none" }, "sha512-2/w+q8jszv9Ww1c+6uJT3OwqhdmGP2/4T17cu8WuwyUuuaCDDJ2ojdyYwZzCxx0GcsZBhzi3HmH+J5pZNXnd+Q=="], + + "@rollup/rollup-linux-ppc64-gnu": ["@rollup/rollup-linux-ppc64-gnu@4.60.2", "", { "os": "linux", "cpu": "ppc64" }, "sha512-11+aL5vKheYgczxtPVVRhdptAM2H7fcDR5Gw4/bTcteuZBlH4oP9f5s9zYO9aGZvoGeBpqXI/9TZZihZ609wKw=="], + + "@rollup/rollup-linux-ppc64-musl": ["@rollup/rollup-linux-ppc64-musl@4.60.2", "", { "os": "linux", "cpu": "ppc64" }, "sha512-i16fokAGK46IVZuV8LIIwMdtqhin9hfYkCh8pf8iC3QU3LpwL+1FSFGej+O7l3E/AoknL6Dclh2oTdnRMpTzFQ=="], + + "@rollup/rollup-linux-riscv64-gnu": ["@rollup/rollup-linux-riscv64-gnu@4.60.2", "", { "os": "linux", "cpu": "none" }, "sha512-49FkKS6RGQoriDSK/6E2GkAsAuU5kETFCh7pG4yD/ylj9rKhTmO3elsnmBvRD4PgJPds5W2PkhC82aVwmUcJ7A=="], + + "@rollup/rollup-linux-riscv64-musl": ["@rollup/rollup-linux-riscv64-musl@4.60.2", "", { "os": "linux", "cpu": "none" }, "sha512-mjYNkHPfGpUR00DuM1ZZIgs64Hpf4bWcz9Z41+4Q+pgDx73UwWdAYyf6EG/lRFldmdHHzgrYyge5akFUW0D3mQ=="], + + "@rollup/rollup-linux-s390x-gnu": ["@rollup/rollup-linux-s390x-gnu@4.60.2", "", { "os": "linux", "cpu": "s390x" }, "sha512-ALyvJz965BQk8E9Al/JDKKDLH2kfKFLTGMlgkAbbYtZuJt9LU8DW3ZoDMCtQpXAltZxwBHevXz5u+gf0yA0YoA=="], + + "@rollup/rollup-linux-x64-gnu": ["@rollup/rollup-linux-x64-gnu@4.60.2", "", { "os": "linux", "cpu": "x64" }, "sha512-UQjrkIdWrKI626Du8lCQ6MJp/6V1LAo2bOK9OTu4mSn8GGXIkPXk/Vsp4bLHCd9Z9Iz2OTEaokUE90VweJgIYQ=="], + + "@rollup/rollup-linux-x64-musl": ["@rollup/rollup-linux-x64-musl@4.60.2", "", { "os": "linux", "cpu": "x64" }, "sha512-bTsRGj6VlSdn/XD4CGyzMnzaBs9bsRxy79eTqTCBsA8TMIEky7qg48aPkvJvFe1HyzQ5oMZdg7AnVlWQSKLTnw=="], + + "@rollup/rollup-openbsd-x64": ["@rollup/rollup-openbsd-x64@4.60.2", "", { "os": "openbsd", "cpu": "x64" }, "sha512-6d4Z3534xitaA1FcMWP7mQPq5zGwBmGbhphh2DwaA1aNIXUu3KTOfwrWpbwI4/Gr0uANo7NTtaykFyO2hPuFLg=="], + + "@rollup/rollup-openharmony-arm64": ["@rollup/rollup-openharmony-arm64@4.60.2", "", { "os": "none", "cpu": "arm64" }, "sha512-NetAg5iO2uN7eB8zE5qrZ3CSil+7IJt4WDFLcC75Ymywq1VZVD6qJ6EvNLjZ3rEm6gB7XW5JdT60c6MN35Z85Q=="], + + "@rollup/rollup-win32-arm64-msvc": ["@rollup/rollup-win32-arm64-msvc@4.60.2", "", { "os": "win32", "cpu": "arm64" }, "sha512-NCYhOotpgWZ5kdxCZsv6Iudx0wX8980Q/oW4pNFNihpBKsDbEA1zpkfxJGC0yugsUuyDZ7gL37dbzwhR0VI7pQ=="], + + "@rollup/rollup-win32-ia32-msvc": ["@rollup/rollup-win32-ia32-msvc@4.60.2", "", { "os": "win32", "cpu": "ia32" }, "sha512-RXsaOqXxfoUBQoOgvmmijVxJnW2IGB0eoMO7F8FAjaj0UTywUO/luSqimWBJn04WNgUkeNhh7fs7pESXajWmkg=="], + + "@rollup/rollup-win32-x64-gnu": ["@rollup/rollup-win32-x64-gnu@4.60.2", "", { "os": "win32", "cpu": "x64" }, "sha512-qdAzEULD+/hzObedtmV6iBpdL5TIbKVztGiK7O3/KYSf+HIzU257+MX1EXJcyIiDbMAqmbwaufcYPvyRryeZtA=="], + + "@rollup/rollup-win32-x64-msvc": ["@rollup/rollup-win32-x64-msvc@4.60.2", "", { "os": "win32", "cpu": "x64" }, "sha512-Nd/SgG27WoA9e+/TdK74KnHz852TLa94ovOYySo/yMPuTmpckK/jIF2jSwS3g7ELSKXK13/cVdmg1Z/DaCWKxA=="], + + "@shikijs/core": ["@shikijs/core@2.5.0", "", { "dependencies": { "@shikijs/engine-javascript": "2.5.0", "@shikijs/engine-oniguruma": "2.5.0", "@shikijs/types": "2.5.0", "@shikijs/vscode-textmate": "^10.0.2", "@types/hast": "^3.0.4", "hast-util-to-html": "^9.0.4" } }, "sha512-uu/8RExTKtavlpH7XqnVYBrfBkUc20ngXiX9NSrBhOVZYv/7XQRKUyhtkeflY5QsxC0GbJThCerruZfsUaSldg=="], + + "@shikijs/engine-javascript": ["@shikijs/engine-javascript@2.5.0", "", { "dependencies": { "@shikijs/types": "2.5.0", "@shikijs/vscode-textmate": "^10.0.2", "oniguruma-to-es": "^3.1.0" } }, "sha512-VjnOpnQf8WuCEZtNUdjjwGUbtAVKuZkVQ/5cHy/tojVVRIRtlWMYVjyWhxOmIq05AlSOv72z7hRNRGVBgQOl0w=="], + + "@shikijs/engine-oniguruma": ["@shikijs/engine-oniguruma@2.5.0", "", { "dependencies": { "@shikijs/types": "2.5.0", "@shikijs/vscode-textmate": "^10.0.2" } }, "sha512-pGd1wRATzbo/uatrCIILlAdFVKdxImWJGQ5rFiB5VZi2ve5xj3Ax9jny8QvkaV93btQEwR/rSz5ERFpC5mKNIw=="], + + "@shikijs/langs": ["@shikijs/langs@2.5.0", "", { "dependencies": { "@shikijs/types": "2.5.0" } }, "sha512-Qfrrt5OsNH5R+5tJ/3uYBBZv3SuGmnRPejV9IlIbFH3HTGLDlkqgHymAlzklVmKBjAaVmkPkyikAV/sQ1wSL+w=="], + + "@shikijs/themes": ["@shikijs/themes@2.5.0", "", { "dependencies": { "@shikijs/types": "2.5.0" } }, "sha512-wGrk+R8tJnO0VMzmUExHR+QdSaPUl/NKs+a4cQQRWyoc3YFbUzuLEi/KWK1hj+8BfHRKm2jNhhJck1dfstJpiw=="], + + "@shikijs/transformers": ["@shikijs/transformers@2.5.0", "", { "dependencies": { "@shikijs/core": "2.5.0", "@shikijs/types": "2.5.0" } }, "sha512-SI494W5X60CaUwgi8u4q4m4s3YAFSxln3tzNjOSYqq54wlVgz0/NbbXEb3mdLbqMBztcmS7bVTaEd2w0qMmfeg=="], + + "@shikijs/types": ["@shikijs/types@2.5.0", "", { "dependencies": { "@shikijs/vscode-textmate": "^10.0.2", "@types/hast": "^3.0.4" } }, "sha512-ygl5yhxki9ZLNuNpPitBWvcy9fsSKKaRuO4BAlMyagszQidxcpLAr0qiW/q43DtSIDxO6hEbtYLiFZNXO/hdGw=="], + + "@shikijs/vscode-textmate": ["@shikijs/vscode-textmate@10.0.2", "", {}, "sha512-83yeghZ2xxin3Nj8z1NMd/NCuca+gsYXswywDy5bHvwlWL8tpTQmzGeUuHd9FC3E/SBEMvzJRwWEOz5gGes9Qg=="], + + "@types/estree": ["@types/estree@1.0.8", "", {}, "sha512-dWHzHa2WqEXI/O1E9OjrocMTKJl2mSrEolh1Iomrv6U+JuNwaHXsXx9bLu5gG7BUWFIN0skIQJQ/L1rIex4X6w=="], + + "@types/hast": ["@types/hast@3.0.4", "", { "dependencies": { "@types/unist": "*" } }, "sha512-WPs+bbQw5aCj+x6laNGWLH3wviHtoCv/P3+otBhbOhJgG8qtpdAMlTCxLtsTWA7LH1Oh/bFCHsBn0TPS5m30EQ=="], + + "@types/linkify-it": ["@types/linkify-it@5.0.0", "", {}, "sha512-sVDA58zAw4eWAffKOaQH5/5j3XeayukzDk+ewSsnv3p4yJEZHCCzMDiZM8e0OUrRvmpGZ85jf4yDHkHsgBNr9Q=="], + + "@types/markdown-it": ["@types/markdown-it@14.1.2", "", { "dependencies": { "@types/linkify-it": "^5", "@types/mdurl": "^2" } }, "sha512-promo4eFwuiW+TfGxhi+0x3czqTYJkG8qB17ZUJiVF10Xm7NLVRSLUsfRTU/6h1e24VvRnXCx+hG7li58lkzog=="], + + "@types/mdast": ["@types/mdast@4.0.4", "", { "dependencies": { "@types/unist": "*" } }, "sha512-kGaNbPh1k7AFzgpud/gMdvIm5xuECykRR+JnWKQno9TAXVa6WIVCGTPvYGekIDL4uwCZQSYbUxNBSb1aUo79oA=="], + + "@types/mdurl": ["@types/mdurl@2.0.0", "", {}, "sha512-RGdgjQUZba5p6QEFAVx2OGb8rQDL/cPRG7GiedRzMcJ1tYnUANBncjbSB1NRGwbvjcPeikRABz2nshyPk1bhWg=="], + + "@types/unist": ["@types/unist@3.0.3", "", {}, "sha512-ko/gIFJRv177XgZsZcBwnqJN5x/Gien8qNOn0D5bQU/zAzVf9Zt3BlcUiLqhV9y4ARk0GbT3tnUiPNgnTXzc/Q=="], + + "@types/web-bluetooth": ["@types/web-bluetooth@0.0.21", "", {}, "sha512-oIQLCGWtcFZy2JW77j9k8nHzAOpqMHLQejDA48XXMWH6tjCQHz5RCFz1bzsmROyL6PUm+LLnUiI4BCn221inxA=="], + + "@ungap/structured-clone": ["@ungap/structured-clone@1.3.0", "", {}, "sha512-WmoN8qaIAo7WTYWbAZuG8PYEhn5fkz7dZrqTBZ7dtt//lL2Gwms1IcnQ5yHqjDfX8Ft5j4YzDM23f87zBfDe9g=="], + + "@vitejs/plugin-vue": ["@vitejs/plugin-vue@5.2.4", "", { "peerDependencies": { "vite": "^5.0.0 || ^6.0.0", "vue": "^3.2.25" } }, "sha512-7Yx/SXSOcQq5HiiV3orevHUFn+pmMB4cgbEkDYgnkUWb0WfeQ/wa2yFv6D5ICiCQOVpjA7vYDXrC7AGO8yjDHA=="], + + "@vue/compiler-core": ["@vue/compiler-core@3.5.32", "", { "dependencies": { "@babel/parser": "^7.29.2", "@vue/shared": "3.5.32", "entities": "^7.0.1", "estree-walker": "^2.0.2", "source-map-js": "^1.2.1" } }, "sha512-4x74Tbtqnda8s/NSD6e1Dr5p1c8HdMU5RWSjMSUzb8RTcUQqevDCxVAitcLBKT+ie3o0Dl9crc/S/opJM7qBGQ=="], + + "@vue/compiler-dom": ["@vue/compiler-dom@3.5.32", "", { "dependencies": { "@vue/compiler-core": "3.5.32", "@vue/shared": "3.5.32" } }, "sha512-ybHAu70NtiEI1fvAUz3oXZqkUYEe5J98GjMDpTGl5iHb0T15wQYLR4wE3h9xfuTNA+Cm2f4czfe8B4s+CCH57Q=="], + + "@vue/compiler-sfc": ["@vue/compiler-sfc@3.5.32", "", { "dependencies": { "@babel/parser": "^7.29.2", "@vue/compiler-core": "3.5.32", "@vue/compiler-dom": "3.5.32", "@vue/compiler-ssr": "3.5.32", "@vue/shared": "3.5.32", "estree-walker": "^2.0.2", "magic-string": "^0.30.21", "postcss": "^8.5.8", "source-map-js": "^1.2.1" } }, "sha512-8UYUYo71cP/0YHMO814TRZlPuUUw3oifHuMR7Wp9SNoRSrxRQnhMLNlCeaODNn6kNTJsjFoQ/kqIj4qGvya4Xg=="], + + "@vue/compiler-ssr": ["@vue/compiler-ssr@3.5.32", "", { "dependencies": { "@vue/compiler-dom": "3.5.32", "@vue/shared": "3.5.32" } }, "sha512-Gp4gTs22T3DgRotZ8aA/6m2jMR+GMztvBXUBEUOYOcST+giyGWJ4WvFd7QLHBkzTxkfOt8IELKNdpzITLbA2rw=="], + + "@vue/devtools-api": ["@vue/devtools-api@7.7.9", "", { "dependencies": { "@vue/devtools-kit": "^7.7.9" } }, "sha512-kIE8wvwlcZ6TJTbNeU2HQNtaxLx3a84aotTITUuL/4bzfPxzajGBOoqjMhwZJ8L9qFYDU/lAYMEEm11dnZOD6g=="], + + "@vue/devtools-kit": ["@vue/devtools-kit@7.7.9", "", { "dependencies": { "@vue/devtools-shared": "^7.7.9", "birpc": "^2.3.0", "hookable": "^5.5.3", "mitt": "^3.0.1", "perfect-debounce": "^1.0.0", "speakingurl": "^14.0.1", "superjson": "^2.2.2" } }, "sha512-PyQ6odHSgiDVd4hnTP+aDk2X4gl2HmLDfiyEnn3/oV+ckFDuswRs4IbBT7vacMuGdwY/XemxBoh302ctbsptuA=="], + + "@vue/devtools-shared": ["@vue/devtools-shared@7.7.9", "", { "dependencies": { "rfdc": "^1.4.1" } }, "sha512-iWAb0v2WYf0QWmxCGy0seZNDPdO3Sp5+u78ORnyeonS6MT4PC7VPrryX2BpMJrwlDeaZ6BD4vP4XKjK0SZqaeA=="], + + "@vue/reactivity": ["@vue/reactivity@3.5.32", "", { "dependencies": { "@vue/shared": "3.5.32" } }, "sha512-/ORasxSGvZ6MN5gc+uE364SxFdJ0+WqVG0CENXaGW58TOCdrAW76WWaplDtECeS1qphvtBZtR+3/o1g1zL4xPQ=="], + + "@vue/runtime-core": ["@vue/runtime-core@3.5.32", "", { "dependencies": { "@vue/reactivity": "3.5.32", "@vue/shared": "3.5.32" } }, "sha512-pDrXCejn4UpFDFmMd27AcJEbHaLemaE5o4pbb7sLk79SRIhc6/t34BQA7SGNgYtbMnvbF/HHOftYBgFJtUoJUQ=="], + + "@vue/runtime-dom": ["@vue/runtime-dom@3.5.32", "", { "dependencies": { "@vue/reactivity": "3.5.32", "@vue/runtime-core": "3.5.32", "@vue/shared": "3.5.32", "csstype": "^3.2.3" } }, "sha512-1CDVv7tv/IV13V8Nip1k/aaObVbWqRlVCVezTwx3K07p7Vxossp5JU1dcPNhJk3w347gonIUT9jQOGutyJrSVQ=="], + + "@vue/server-renderer": ["@vue/server-renderer@3.5.32", "", { "dependencies": { "@vue/compiler-ssr": "3.5.32", "@vue/shared": "3.5.32" }, "peerDependencies": { "vue": "3.5.32" } }, "sha512-IOjm2+JQwRFS7W28HNuJeXQle9KdZbODFY7hFGVtnnghF51ta20EWAZJHX+zLGtsHhaU6uC9BGPV52KVpYryMQ=="], + + "@vue/shared": ["@vue/shared@3.5.32", "", {}, "sha512-ksNyrmRQzWJJ8n3cRDuSF7zNNontuJg1YHnmWRJd2AMu8Ij2bqwiiri2lH5rHtYPZjj4STkNcgcmiQqlOjiYGg=="], + + "@vueuse/core": ["@vueuse/core@12.8.2", "", { "dependencies": { "@types/web-bluetooth": "^0.0.21", "@vueuse/metadata": "12.8.2", "@vueuse/shared": "12.8.2", "vue": "^3.5.13" } }, "sha512-HbvCmZdzAu3VGi/pWYm5Ut+Kd9mn1ZHnn4L5G8kOQTPs/IwIAmJoBrmYk2ckLArgMXZj0AW3n5CAejLUO+PhdQ=="], + + "@vueuse/integrations": ["@vueuse/integrations@12.8.2", "", { "dependencies": { "@vueuse/core": "12.8.2", "@vueuse/shared": "12.8.2", "vue": "^3.5.13" }, "peerDependencies": { "async-validator": "^4", "axios": "^1", "change-case": "^5", "drauu": "^0.4", "focus-trap": "^7", "fuse.js": "^7", "idb-keyval": "^6", "jwt-decode": "^4", "nprogress": "^0.2", "qrcode": "^1.5", "sortablejs": "^1", "universal-cookie": "^7" }, "optionalPeers": ["async-validator", "axios", "change-case", "drauu", "focus-trap", "fuse.js", "idb-keyval", "jwt-decode", "nprogress", "qrcode", "sortablejs", "universal-cookie"] }, "sha512-fbGYivgK5uBTRt7p5F3zy6VrETlV9RtZjBqd1/HxGdjdckBgBM4ugP8LHpjolqTj14TXTxSK1ZfgPbHYyGuH7g=="], + + "@vueuse/metadata": ["@vueuse/metadata@12.8.2", "", {}, "sha512-rAyLGEuoBJ/Il5AmFHiziCPdQzRt88VxR+Y/A/QhJ1EWtWqPBBAxTAFaSkviwEuOEZNtW8pvkPgoCZQ+HxqW1A=="], + + "@vueuse/shared": ["@vueuse/shared@12.8.2", "", { "dependencies": { "vue": "^3.5.13" } }, "sha512-dznP38YzxZoNloI0qpEfpkms8knDtaoQ6Y/sfS0L7Yki4zh40LFHEhur0odJC6xTHG5dxWVPiUWBXn+wCG2s5w=="], + + "algoliasearch": ["algoliasearch@5.50.2", "", { "dependencies": { "@algolia/abtesting": "1.16.2", "@algolia/client-abtesting": "5.50.2", "@algolia/client-analytics": "5.50.2", "@algolia/client-common": "5.50.2", "@algolia/client-insights": "5.50.2", "@algolia/client-personalization": "5.50.2", "@algolia/client-query-suggestions": "5.50.2", "@algolia/client-search": "5.50.2", "@algolia/ingestion": "1.50.2", "@algolia/monitoring": "1.50.2", "@algolia/recommend": "5.50.2", "@algolia/requester-browser-xhr": "5.50.2", "@algolia/requester-fetch": "5.50.2", "@algolia/requester-node-http": "5.50.2" } }, "sha512-Tfp26yoNWurUjfgK4GOrVJQhSNXu9tJtHfFFNosgT2YClG+vPyUjX/gbC8rG39qLncnZg8Fj34iarQWpMkqefw=="], + + "birpc": ["birpc@2.9.0", "", {}, "sha512-KrayHS5pBi69Xi9JmvoqrIgYGDkD6mcSe/i6YKi3w5kekCLzrX4+nawcXqrj2tIp50Kw/mT/s3p+GVK0A0sKxw=="], + + "ccount": ["ccount@2.0.1", "", {}, "sha512-eyrF0jiFpY+3drT6383f1qhkbGsLSifNAjA61IUjZjmLCWjItY6LB9ft9YhoDgwfmclB2zhu51Lc7+95b8NRAg=="], + + "character-entities-html4": ["character-entities-html4@2.1.0", "", {}, "sha512-1v7fgQRj6hnSwFpq1Eu0ynr/CDEw0rXo2B61qXrLNdHZmPKgb7fqS1a2JwF0rISo9q77jDI8VMEHoApn8qDoZA=="], + + "character-entities-legacy": ["character-entities-legacy@3.0.0", "", {}, "sha512-RpPp0asT/6ufRm//AJVwpViZbGM/MkjQFxJccQRHmISF/22NBtsHqAWmL+/pmkPWoIUJdWyeVleTl1wydHATVQ=="], + + "comma-separated-tokens": ["comma-separated-tokens@2.0.3", "", {}, "sha512-Fu4hJdvzeylCfQPp9SGWidpzrMs7tTrlu6Vb8XGaRGck8QSNZJJp538Wrb60Lax4fPwR64ViY468OIUTbRlGZg=="], + + "copy-anything": ["copy-anything@4.0.5", "", { "dependencies": { "is-what": "^5.2.0" } }, "sha512-7Vv6asjS4gMOuILabD3l739tsaxFQmC+a7pLZm02zyvs8p977bL3zEgq3yDk5rn9B0PbYgIv++jmHcuUab4RhA=="], + + "csstype": ["csstype@3.2.3", "", {}, "sha512-z1HGKcYy2xA8AGQfwrn0PAy+PB7X/GSj3UVJW9qKyn43xWa+gl5nXmU4qqLMRzWVLFC8KusUX8T/0kCiOYpAIQ=="], + + "dequal": ["dequal@2.0.3", "", {}, "sha512-0je+qPKHEMohvfRTCEo3CrPG6cAzAYgmzKyxRiYSSDkS6eGJdyVJm7WaYA5ECaAD9wLB2T4EEeymA5aFVcYXCA=="], + + "devlop": ["devlop@1.1.0", "", { "dependencies": { "dequal": "^2.0.0" } }, "sha512-RWmIqhcFf1lRYBvNmr7qTNuyCt/7/ns2jbpp1+PalgE/rDQcBT0fioSMUpJ93irlUhC5hrg4cYqe6U+0ImW0rA=="], + + "emoji-regex-xs": ["emoji-regex-xs@1.0.0", "", {}, "sha512-LRlerrMYoIDrT6jgpeZ2YYl/L8EulRTt5hQcYjy5AInh7HWXKimpqx68aknBFpGL2+/IcogTcaydJEgaTmOpDg=="], + + "entities": ["entities@7.0.1", "", {}, "sha512-TWrgLOFUQTH994YUyl1yT4uyavY5nNB5muff+RtWaqNVCAK408b5ZnnbNAUEWLTCpum9w6arT70i1XdQ4UeOPA=="], + + "esbuild": ["esbuild@0.21.5", "", { "optionalDependencies": { "@esbuild/aix-ppc64": "0.21.5", "@esbuild/android-arm": "0.21.5", "@esbuild/android-arm64": "0.21.5", "@esbuild/android-x64": "0.21.5", "@esbuild/darwin-arm64": "0.21.5", "@esbuild/darwin-x64": "0.21.5", "@esbuild/freebsd-arm64": "0.21.5", "@esbuild/freebsd-x64": "0.21.5", "@esbuild/linux-arm": "0.21.5", "@esbuild/linux-arm64": "0.21.5", "@esbuild/linux-ia32": "0.21.5", "@esbuild/linux-loong64": "0.21.5", "@esbuild/linux-mips64el": "0.21.5", "@esbuild/linux-ppc64": "0.21.5", "@esbuild/linux-riscv64": "0.21.5", "@esbuild/linux-s390x": "0.21.5", "@esbuild/linux-x64": "0.21.5", "@esbuild/netbsd-x64": "0.21.5", "@esbuild/openbsd-x64": "0.21.5", "@esbuild/sunos-x64": "0.21.5", "@esbuild/win32-arm64": "0.21.5", "@esbuild/win32-ia32": "0.21.5", "@esbuild/win32-x64": "0.21.5" }, "bin": { "esbuild": "bin/esbuild" } }, "sha512-mg3OPMV4hXywwpoDxu3Qda5xCKQi+vCTZq8S9J/EpkhB2HzKXq4SNFZE3+NK93JYxc8VMSep+lOUSC/RVKaBqw=="], + + "estree-walker": ["estree-walker@2.0.2", "", {}, "sha512-Rfkk/Mp/DL7JVje3u18FxFujQlTNR2q6QfMSMB7AvCBx91NGj/ba3kCfza0f6dVDbw7YlRf/nDrn7pQrCCyQ/w=="], + + "focus-trap": ["focus-trap@7.8.0", "", { "dependencies": { "tabbable": "^6.4.0" } }, "sha512-/yNdlIkpWbM0ptxno3ONTuf+2g318kh2ez3KSeZN5dZ8YC6AAmgeWz+GasYYiBJPFaYcSAPeu4GfhUaChzIJXA=="], + + "fsevents": ["fsevents@2.3.3", "", { "os": "darwin" }, "sha512-5xoDfX+fL7faATnagmWPpbFtwh/R77WmMMqqHGS65C3vvB0YHrgF+B1YmZ3441tMj5n63k0212XNoJwzlhffQw=="], + + "hast-util-to-html": ["hast-util-to-html@9.0.5", "", { "dependencies": { "@types/hast": "^3.0.0", "@types/unist": "^3.0.0", "ccount": "^2.0.0", "comma-separated-tokens": "^2.0.0", "hast-util-whitespace": "^3.0.0", "html-void-elements": "^3.0.0", "mdast-util-to-hast": "^13.0.0", "property-information": "^7.0.0", "space-separated-tokens": "^2.0.0", "stringify-entities": "^4.0.0", "zwitch": "^2.0.4" } }, "sha512-OguPdidb+fbHQSU4Q4ZiLKnzWo8Wwsf5bZfbvu7//a9oTYoqD/fWpe96NuHkoS9h0ccGOTe0C4NGXdtS0iObOw=="], + + "hast-util-whitespace": ["hast-util-whitespace@3.0.0", "", { "dependencies": { "@types/hast": "^3.0.0" } }, "sha512-88JUN06ipLwsnv+dVn+OIYOvAuvBMy/Qoi6O7mQHxdPXpjy+Cd6xRkWwux7DKO+4sYILtLBRIKgsdpS2gQc7qw=="], + + "hookable": ["hookable@5.5.3", "", {}, "sha512-Yc+BQe8SvoXH1643Qez1zqLRmbA5rCL+sSmk6TVos0LWVfNIB7PGncdlId77WzLGSIB5KaWgTaNTs2lNVEI6VQ=="], + + "html-void-elements": ["html-void-elements@3.0.0", "", {}, "sha512-bEqo66MRXsUGxWHV5IP0PUiAWwoEjba4VCzg0LjFJBpchPaTfyfCKTG6bc5F8ucKec3q5y6qOdGyYTSBEvhCrg=="], + + "is-what": ["is-what@5.5.0", "", {}, "sha512-oG7cgbmg5kLYae2N5IVd3jm2s+vldjxJzK1pcu9LfpGuQ93MQSzo0okvRna+7y5ifrD+20FE8FvjusyGaz14fw=="], + + "magic-string": ["magic-string@0.30.21", "", { "dependencies": { "@jridgewell/sourcemap-codec": "^1.5.5" } }, "sha512-vd2F4YUyEXKGcLHoq+TEyCjxueSeHnFxyyjNp80yg0XV4vUhnDer/lvvlqM/arB5bXQN5K2/3oinyCRyx8T2CQ=="], + + "mark.js": ["mark.js@8.11.1", "", {}, "sha512-1I+1qpDt4idfgLQG+BNWmrqku+7/2bi5nLf4YwF8y8zXvmfiTBY3PV3ZibfrjBueCByROpuBjLLFCajqkgYoLQ=="], + + "mdast-util-to-hast": ["mdast-util-to-hast@13.2.1", "", { "dependencies": { "@types/hast": "^3.0.0", "@types/mdast": "^4.0.0", "@ungap/structured-clone": "^1.0.0", "devlop": "^1.0.0", "micromark-util-sanitize-uri": "^2.0.0", "trim-lines": "^3.0.0", "unist-util-position": "^5.0.0", "unist-util-visit": "^5.0.0", "vfile": "^6.0.0" } }, "sha512-cctsq2wp5vTsLIcaymblUriiTcZd0CwWtCbLvrOzYCDZoWyMNV8sZ7krj09FSnsiJi3WVsHLM4k6Dq/yaPyCXA=="], + + "micromark-util-character": ["micromark-util-character@2.1.1", "", { "dependencies": { "micromark-util-symbol": "^2.0.0", "micromark-util-types": "^2.0.0" } }, "sha512-wv8tdUTJ3thSFFFJKtpYKOYiGP2+v96Hvk4Tu8KpCAsTMs6yi+nVmGh1syvSCsaxz45J6Jbw+9DD6g97+NV67Q=="], + + "micromark-util-encode": ["micromark-util-encode@2.0.1", "", {}, "sha512-c3cVx2y4KqUnwopcO9b/SCdo2O67LwJJ/UyqGfbigahfegL9myoEFoDYZgkT7f36T0bLrM9hZTAaAyH+PCAXjw=="], + + "micromark-util-sanitize-uri": ["micromark-util-sanitize-uri@2.0.1", "", { "dependencies": { "micromark-util-character": "^2.0.0", "micromark-util-encode": "^2.0.0", "micromark-util-symbol": "^2.0.0" } }, "sha512-9N9IomZ/YuGGZZmQec1MbgxtlgougxTodVwDzzEouPKo3qFWvymFHWcnDi2vzV1ff6kas9ucW+o3yzJK9YB1AQ=="], + + "micromark-util-symbol": ["micromark-util-symbol@2.0.1", "", {}, "sha512-vs5t8Apaud9N28kgCrRUdEed4UJ+wWNvicHLPxCa9ENlYuAY31M0ETy5y1vA33YoNPDFTghEbnh6efaE8h4x0Q=="], + + "micromark-util-types": ["micromark-util-types@2.0.2", "", {}, "sha512-Yw0ECSpJoViF1qTU4DC6NwtC4aWGt1EkzaQB8KPPyCRR8z9TWeV0HbEFGTO+ZY1wB22zmxnJqhPyTpOVCpeHTA=="], + + "minisearch": ["minisearch@7.2.0", "", {}, "sha512-dqT2XBYUOZOiC5t2HRnwADjhNS2cecp9u+TJRiJ1Qp/f5qjkeT5APcGPjHw+bz89Ms8Jp+cG4AlE+QZ/QnDglg=="], + + "mitt": ["mitt@3.0.1", "", {}, "sha512-vKivATfr97l2/QBCYAkXYDbrIWPM2IIKEl7YPhjCvKlG3kE2gm+uBo6nEXK3M5/Ffh/FLpKExzOQ3JJoJGFKBw=="], + + "nanoid": ["nanoid@3.3.11", "", { "bin": { "nanoid": "bin/nanoid.cjs" } }, "sha512-N8SpfPUnUp1bK+PMYW8qSWdl9U+wwNWI4QKxOYDy9JAro3WMX7p2OeVRF9v+347pnakNevPmiHhNmZ2HbFA76w=="], + + "oniguruma-to-es": ["oniguruma-to-es@3.1.1", "", { "dependencies": { "emoji-regex-xs": "^1.0.0", "regex": "^6.0.1", "regex-recursion": "^6.0.2" } }, "sha512-bUH8SDvPkH3ho3dvwJwfonjlQ4R80vjyvrU8YpxuROddv55vAEJrTuCuCVUhhsHbtlD9tGGbaNApGQckXhS8iQ=="], + + "perfect-debounce": ["perfect-debounce@1.0.0", "", {}, "sha512-xCy9V055GLEqoFaHoC1SoLIaLmWctgCUaBaWxDZ7/Zx4CTyX7cJQLJOok/orfjZAh9kEYpjJa4d0KcJmCbctZA=="], + + "picocolors": ["picocolors@1.1.1", "", {}, "sha512-xceH2snhtb5M9liqDsmEw56le376mTZkEX/jEb/RxNFyegNul7eNslCXP9FDj/Lcu0X8KEyMceP2ntpaHrDEVA=="], + + "postcss": ["postcss@8.5.10", "", { "dependencies": { "nanoid": "^3.3.11", "picocolors": "^1.1.1", "source-map-js": "^1.2.1" } }, "sha512-pMMHxBOZKFU6HgAZ4eyGnwXF/EvPGGqUr0MnZ5+99485wwW41kW91A4LOGxSHhgugZmSChL5AlElNdwlNgcnLQ=="], + + "preact": ["preact@10.29.1", "", {}, "sha512-gQCLc/vWroE8lIpleXtdJhTFDogTdZG9AjMUpVkDf2iTCNwYNWA+u16dL41TqUDJO4gm2IgrcMv3uTpjd4Pwmg=="], + + "property-information": ["property-information@7.1.0", "", {}, "sha512-TwEZ+X+yCJmYfL7TPUOcvBZ4QfoT5YenQiJuX//0th53DE6w0xxLEtfK3iyryQFddXuvkIk51EEgrJQ0WJkOmQ=="], + + "regex": ["regex@6.1.0", "", { "dependencies": { "regex-utilities": "^2.3.0" } }, "sha512-6VwtthbV4o/7+OaAF9I5L5V3llLEsoPyq9P1JVXkedTP33c7MfCG0/5NOPcSJn0TzXcG9YUrR0gQSWioew3LDg=="], + + "regex-recursion": ["regex-recursion@6.0.2", "", { "dependencies": { "regex-utilities": "^2.3.0" } }, "sha512-0YCaSCq2VRIebiaUviZNs0cBz1kg5kVS2UKUfNIx8YVs1cN3AV7NTctO5FOKBA+UT2BPJIWZauYHPqJODG50cg=="], + + "regex-utilities": ["regex-utilities@2.3.0", "", {}, "sha512-8VhliFJAWRaUiVvREIiW2NXXTmHs4vMNnSzuJVhscgmGav3g9VDxLrQndI3dZZVVdp0ZO/5v0xmX516/7M9cng=="], + + "rfdc": ["rfdc@1.4.1", "", {}, "sha512-q1b3N5QkRUWUl7iyylaaj3kOpIT0N2i9MqIEQXP73GVsN9cw3fdx8X63cEmWhJGi2PPCF23Ijp7ktmd39rawIA=="], + + "rollup": ["rollup@4.60.2", "", { "dependencies": { "@types/estree": "1.0.8" }, "optionalDependencies": { "@rollup/rollup-android-arm-eabi": "4.60.2", "@rollup/rollup-android-arm64": "4.60.2", "@rollup/rollup-darwin-arm64": "4.60.2", "@rollup/rollup-darwin-x64": "4.60.2", "@rollup/rollup-freebsd-arm64": "4.60.2", "@rollup/rollup-freebsd-x64": "4.60.2", "@rollup/rollup-linux-arm-gnueabihf": "4.60.2", "@rollup/rollup-linux-arm-musleabihf": "4.60.2", "@rollup/rollup-linux-arm64-gnu": "4.60.2", "@rollup/rollup-linux-arm64-musl": "4.60.2", "@rollup/rollup-linux-loong64-gnu": "4.60.2", "@rollup/rollup-linux-loong64-musl": "4.60.2", "@rollup/rollup-linux-ppc64-gnu": "4.60.2", "@rollup/rollup-linux-ppc64-musl": "4.60.2", "@rollup/rollup-linux-riscv64-gnu": "4.60.2", "@rollup/rollup-linux-riscv64-musl": "4.60.2", "@rollup/rollup-linux-s390x-gnu": "4.60.2", "@rollup/rollup-linux-x64-gnu": "4.60.2", "@rollup/rollup-linux-x64-musl": "4.60.2", "@rollup/rollup-openbsd-x64": "4.60.2", "@rollup/rollup-openharmony-arm64": "4.60.2", "@rollup/rollup-win32-arm64-msvc": "4.60.2", "@rollup/rollup-win32-ia32-msvc": "4.60.2", "@rollup/rollup-win32-x64-gnu": "4.60.2", "@rollup/rollup-win32-x64-msvc": "4.60.2", "fsevents": "~2.3.2" }, "bin": { "rollup": "dist/bin/rollup" } }, "sha512-J9qZyW++QK/09NyN/zeO0dG/1GdGfyp9lV8ajHnRVLfo/uFsbji5mHnDgn/qYdUHyCkM2N+8VyspgZclfAh0eQ=="], + + "search-insights": ["search-insights@2.17.3", "", {}, "sha512-RQPdCYTa8A68uM2jwxoY842xDhvx3E5LFL1LxvxCNMev4o5mLuokczhzjAgGwUZBAmOKZknArSxLKmXtIi2AxQ=="], + + "shiki": ["shiki@2.5.0", "", { "dependencies": { "@shikijs/core": "2.5.0", "@shikijs/engine-javascript": "2.5.0", "@shikijs/engine-oniguruma": "2.5.0", "@shikijs/langs": "2.5.0", "@shikijs/themes": "2.5.0", "@shikijs/types": "2.5.0", "@shikijs/vscode-textmate": "^10.0.2", "@types/hast": "^3.0.4" } }, "sha512-mI//trrsaiCIPsja5CNfsyNOqgAZUb6VpJA+340toL42UpzQlXpwRV9nch69X6gaUxrr9kaOOa6e3y3uAkGFxQ=="], + + "source-map-js": ["source-map-js@1.2.1", "", {}, "sha512-UXWMKhLOwVKb728IUtQPXxfYU+usdybtUrK/8uGE8CQMvrhOpwvzDBwj0QhSL7MQc7vIsISBG8VQ8+IDQxpfQA=="], + + "space-separated-tokens": ["space-separated-tokens@2.0.2", "", {}, "sha512-PEGlAwrG8yXGXRjW32fGbg66JAlOAwbObuqVoJpv/mRgoWDQfgH1wDPvtzWyUSNAXBGSk8h755YDbbcEy3SH2Q=="], + + "speakingurl": ["speakingurl@14.0.1", "", {}, "sha512-1POYv7uv2gXoyGFpBCmpDVSNV74IfsWlDW216UPjbWufNf+bSU6GdbDsxdcxtfwb4xlI3yxzOTKClUosxARYrQ=="], + + "stringify-entities": ["stringify-entities@4.0.4", "", { "dependencies": { "character-entities-html4": "^2.0.0", "character-entities-legacy": "^3.0.0" } }, "sha512-IwfBptatlO+QCJUo19AqvrPNqlVMpW9YEL2LIVY+Rpv2qsjCGxaDLNRgeGsQWJhfItebuJhsGSLjaBbNSQ+ieg=="], + + "superjson": ["superjson@2.2.6", "", { "dependencies": { "copy-anything": "^4" } }, "sha512-H+ue8Zo4vJmV2nRjpx86P35lzwDT3nItnIsocgumgr0hHMQ+ZGq5vrERg9kJBo5AWGmxZDhzDo+WVIJqkB0cGA=="], + + "tabbable": ["tabbable@6.4.0", "", {}, "sha512-05PUHKSNE8ou2dwIxTngl4EzcnsCDZGJ/iCLtDflR/SHB/ny14rXc+qU5P4mG9JkusiV7EivzY9Mhm55AzAvCg=="], + + "trim-lines": ["trim-lines@3.0.1", "", {}, "sha512-kRj8B+YHZCc9kQYdWfJB2/oUl9rA99qbowYYBtr4ui4mZyAQ2JpvVBd/6U2YloATfqBhBTSMhTpgBHtU0Mf3Rg=="], + + "unist-util-is": ["unist-util-is@6.0.1", "", { "dependencies": { "@types/unist": "^3.0.0" } }, "sha512-LsiILbtBETkDz8I9p1dQ0uyRUWuaQzd/cuEeS1hoRSyW5E5XGmTzlwY1OrNzzakGowI9Dr/I8HVaw4hTtnxy8g=="], + + "unist-util-position": ["unist-util-position@5.0.0", "", { "dependencies": { "@types/unist": "^3.0.0" } }, "sha512-fucsC7HjXvkB5R3kTCO7kUjRdrS0BJt3M/FPxmHMBOm8JQi2BsHAHFsy27E0EolP8rp0NzXsJ+jNPyDWvOJZPA=="], + + "unist-util-stringify-position": ["unist-util-stringify-position@4.0.0", "", { "dependencies": { "@types/unist": "^3.0.0" } }, "sha512-0ASV06AAoKCDkS2+xw5RXJywruurpbC4JZSm7nr7MOt1ojAzvyyaO+UxZf18j8FCF6kmzCZKcAgN/yu2gm2XgQ=="], + + "unist-util-visit": ["unist-util-visit@5.1.0", "", { "dependencies": { "@types/unist": "^3.0.0", "unist-util-is": "^6.0.0", "unist-util-visit-parents": "^6.0.0" } }, "sha512-m+vIdyeCOpdr/QeQCu2EzxX/ohgS8KbnPDgFni4dQsfSCtpz8UqDyY5GjRru8PDKuYn7Fq19j1CQ+nJSsGKOzg=="], + + "unist-util-visit-parents": ["unist-util-visit-parents@6.0.2", "", { "dependencies": { "@types/unist": "^3.0.0", "unist-util-is": "^6.0.0" } }, "sha512-goh1s1TBrqSqukSc8wrjwWhL0hiJxgA8m4kFxGlQ+8FYQ3C/m11FcTs4YYem7V664AhHVvgoQLk890Ssdsr2IQ=="], + + "vfile": ["vfile@6.0.3", "", { "dependencies": { "@types/unist": "^3.0.0", "vfile-message": "^4.0.0" } }, "sha512-KzIbH/9tXat2u30jf+smMwFCsno4wHVdNmzFyL+T/L3UGqqk6JKfVqOFOZEpZSHADH1k40ab6NUIXZq422ov3Q=="], + + "vfile-message": ["vfile-message@4.0.3", "", { "dependencies": { "@types/unist": "^3.0.0", "unist-util-stringify-position": "^4.0.0" } }, "sha512-QTHzsGd1EhbZs4AsQ20JX1rC3cOlt/IWJruk893DfLRr57lcnOeMaWG4K0JrRta4mIJZKth2Au3mM3u03/JWKw=="], + + "vite": ["vite@5.4.21", "", { "dependencies": { "esbuild": "^0.21.3", "postcss": "^8.4.43", "rollup": "^4.20.0" }, "optionalDependencies": { "fsevents": "~2.3.3" }, "peerDependencies": { "@types/node": "^18.0.0 || >=20.0.0", "less": "*", "lightningcss": "^1.21.0", "sass": "*", "sass-embedded": "*", "stylus": "*", "sugarss": "*", "terser": "^5.4.0" }, "optionalPeers": ["@types/node", "less", "lightningcss", "sass", "sass-embedded", "stylus", "sugarss", "terser"], "bin": { "vite": "bin/vite.js" } }, "sha512-o5a9xKjbtuhY6Bi5S3+HvbRERmouabWbyUcpXXUA1u+GNUKoROi9byOJ8M0nHbHYHkYICiMlqxkg1KkYmm25Sw=="], + + "vitepress": ["vitepress@1.6.4", "", { "dependencies": { "@docsearch/css": "3.8.2", "@docsearch/js": "3.8.2", "@iconify-json/simple-icons": "^1.2.21", "@shikijs/core": "^2.1.0", "@shikijs/transformers": "^2.1.0", "@shikijs/types": "^2.1.0", "@types/markdown-it": "^14.1.2", "@vitejs/plugin-vue": "^5.2.1", "@vue/devtools-api": "^7.7.0", "@vue/shared": "^3.5.13", "@vueuse/core": "^12.4.0", "@vueuse/integrations": "^12.4.0", "focus-trap": "^7.6.4", "mark.js": "8.11.1", "minisearch": "^7.1.1", "shiki": "^2.1.0", "vite": "^5.4.14", "vue": "^3.5.13" }, "peerDependencies": { "markdown-it-mathjax3": "^4", "postcss": "^8" }, "optionalPeers": ["markdown-it-mathjax3", "postcss"], "bin": { "vitepress": "bin/vitepress.js" } }, "sha512-+2ym1/+0VVrbhNyRoFFesVvBvHAVMZMK0rw60E3X/5349M1GuVdKeazuksqopEdvkKwKGs21Q729jX81/bkBJg=="], + + "vue": ["vue@3.5.32", "", { "dependencies": { "@vue/compiler-dom": "3.5.32", "@vue/compiler-sfc": "3.5.32", "@vue/runtime-dom": "3.5.32", "@vue/server-renderer": "3.5.32", "@vue/shared": "3.5.32" }, "peerDependencies": { "typescript": "*" }, "optionalPeers": ["typescript"] }, "sha512-vM4z4Q9tTafVfMAK7IVzmxg34rSzTFMyIe0UUEijUCkn9+23lj0WRfA83dg7eQZIUlgOSGrkViIaCfqSAUXsMw=="], + + "zwitch": ["zwitch@2.0.4", "", {}, "sha512-bXE4cR/kVZhKZX/RjPEflHaKVhUVl85noU3v6b8apfQEc1x4A+zBxjZ4lN8LqGd6WZ3dl98pY4o717VFmoPp+A=="], + } +} diff --git a/docs/configuration/compatible-parameters.md b/docs/configuration/compatible-parameters.md new file mode 100644 index 0000000..876a50b --- /dev/null +++ b/docs/configuration/compatible-parameters.md @@ -0,0 +1,64 @@ +# Compatible Parameters + +TokenSpeed keeps familiar serving parameter names when the operational meaning +is the same. This makes recipes portable while still documenting +TokenSpeed-specific behavior explicitly. + +## Directly Aligned + +| Parameter | TokenSpeed behavior | +| --- | --- | +| positional `model` | Model path or Hugging Face repo ID. | +| `--model` | Equivalent to positional `model`. | +| `--tokenizer` | Tokenizer path. | +| `--tokenizer-mode` | Tokenizer implementation mode. | +| `--skip-tokenizer-init` | Skip tokenizer initialization. | +| `--load-format` | Weight loading format. | +| `--trust-remote-code` | Allow custom model code from the model repository. | +| `--dtype` | Weight and activation dtype. | +| `--kv-cache-dtype` | KV cache storage dtype. | +| `--quantization` | Weight quantization method. | +| `--quantization-param-path` | KV cache scaling-factor file. | +| `--max-model-len` | Maximum sequence length. | +| `--device` | Device type. TokenSpeed currently serves CUDA. | +| `--served-model-name` | OpenAI-compatible served model name. | +| `--revision` | Model revision. | +| `--download-dir` | Model download directory. | +| `--hf-overrides` | JSON model config overrides. | +| `--host` | HTTP bind host. | +| `--port` | HTTP bind port. | +| `--api-key` | API key for the server. | +| `--chat-template` | Chat template name or path. | +| `--gpu-memory-utilization` | GPU memory fraction used for weights and KV cache. | +| `--max-num-seqs` | Maximum concurrent sequences. | +| `--block-size` | KV cache block size. | +| `--enable-prefix-caching` | Enable prefix cache reuse. | +| `--no-enable-prefix-caching` | Disable prefix cache reuse. | +| `--enforce-eager` | Disable CUDA graph execution. | +| `--max-cudagraph-capture-size` | Largest CUDA graph capture size. | +| `--tensor-parallel-size`, `--tp` | Set attention tensor parallel size. | +| `--data-parallel-size` | Data parallel size. | +| `--enable-expert-parallel` | Enable expert parallelism. | +| `--speculative-config` | JSON speculative decoding config. | +| `--kv-events-config` | JSON KV cache event publisher config; the vLLM-style `enable_kv_cache_events` field is accepted and defaults to ZMQ when enabled. | +| `--tool-call-parser` | OpenAI-compatible tool-call parser. | +| `--reasoning-parser` | Reasoning-output parser. | + +## Similar But Not Identical + +| Recipe parameter | TokenSpeed parameter | Difference | +| --- | --- | --- | +| `--max-num-batched-tokens` | `--chunked-prefill-size` | TokenSpeed uses this as the scheduler per-iteration issue budget. | +| `--max-num-batched-tokens` | `--max-total-tokens` | TokenSpeed uses this for the global token pool size override. | +| `--tensor-parallel-size`, `--tp` | `--attn-tp-size` | The familiar alias maps to attention TP. TokenSpeed can split attention, dense, and MoE TP. | +| `--expert-parallel-size` | `--expert-parallel-size`, `--ep-size` | TokenSpeed supports the familiar name and its existing short form. | +| `--attention-backend` | `--attention-backend` | Name is aligned; available backend values are TokenSpeed-specific. | +| `--moe-backend` | `--moe-backend` | Name is aligned; available backend values are TokenSpeed-specific. | + +## Recipe Translation Notes + +- Use `tokenspeed serve` as the launcher. +- Pass the model path positionally, then keep `--trust-remote-code`, `--max-model-len`, `--kv-cache-dtype`, `--gpu-memory-utilization`, `--max-num-seqs`, `--tensor-parallel-size`, `--reasoning-parser`, and `--tool-call-parser` when the model needs them. +- Review `--max-num-batched-tokens` before copying it. TokenSpeed usually wants `--chunked-prefill-size` for per-iteration scheduling. +- Review backend names. TokenSpeed backends are optimized for its runtime and kernel packages. +- Keep TokenSpeed-specific `--attn-tp-size`, `--moe-tp-size`, `--disaggregation-*`, and `--kvstore-*` only when the deployment needs those features. diff --git a/docs/configuration/server.md b/docs/configuration/server.md new file mode 100644 index 0000000..30bc8c1 --- /dev/null +++ b/docs/configuration/server.md @@ -0,0 +1,222 @@ +# Server Parameters + +This page documents the parameters operators usually set directly. TokenSpeed +uses familiar serving parameter names where the semantics match and keeps +TokenSpeed-specific knobs for runtime features with different meaning. + +For a compact compatibility table, see +[Compatible Parameters](./compatible-parameters.md). + +## Model Loading + +| Parameter | Purpose | +| --- | --- | +| positional `model` | Model path or Hugging Face repo ID. | +| `--model` | Equivalent to positional `model`. | +| `--tokenizer` | Tokenizer path when it differs from the model path. | +| `--tokenizer-mode` | Select tokenizer behavior. `auto` uses fast tokenizers and model-specific hooks when available. | +| `--skip-tokenizer-init` | Skip tokenizer initialization for input-ID-only serving paths. | +| `--load-format` | Weight loading format: `auto`, `pt`, `safetensors`, `npcache`, `dummy`, or `extensible`. | +| `--trust-remote-code` | Allow custom model code from the model repository. | +| `--revision` | Model branch, tag, or commit. | +| `--download-dir` | Hugging Face download/cache directory. | +| `--hf-overrides` | JSON overrides for model configuration values. | + +## Precision And Quantization + +| Parameter | Purpose | +| --- | --- | +| `--dtype` | Model weight and activation dtype. `auto` follows model metadata. | +| `--kv-cache-dtype` | KV cache dtype. Lower precision reduces KV memory and may require scaling factors. | +| `--kv-cache-quant-method` | KV cache quantization method. | +| `--quantization` | Weight quantization mode such as `fp8`, `nvfp4`, `w8a8_fp8`, or `compressed-tensors`. | +| `--quantization-param-path` | JSON file for KV cache scaling factors, commonly needed with FP8 KV cache. | + +## API Surface + +| Parameter | Purpose | +| --- | --- | +| `--host` | HTTP bind host. | +| `--port` | HTTP bind port. | +| `--served-model-name` | Model name returned by the OpenAI-compatible API. | +| `--api-key` | API key required by the server. | +| `--chat-template` | Built-in chat template name or template file path (handled by the smg gateway). | +| `--stream-interval` | Streaming buffer interval in generated tokens. Smaller values stream more frequently. | +| `--stream-output` | Return generated text as disjoint streaming segments. | + +## Scheduler And Memory + +| Parameter | Purpose | +| --- | --- | +| `--max-model-len` | Maximum sequence length. If omitted, TokenSpeed uses the model config. | +| `--gpu-memory-utilization` | Fraction of GPU memory used for model weights and KV cache. Lower it to leave headroom. | +| `--max-num-seqs` | Maximum number of active sequences the scheduler may process concurrently. | +| `--chunked-prefill-size` | Token budget the scheduler may issue in one iteration. Defaults to `8192`. Set `-1` to disable chunked prefill. | +| `--max-prefill-tokens` | Prefill token budget used when chunked prefill is disabled. Defaults to `8192`. | +| `--max-total-tokens` | Override the automatically calculated token pool size. | +| `--block-size` | KV cache block size. | +| `--enable-prefix-caching` / `--no-enable-prefix-caching` | Enable or disable prefix cache reuse. | +| `--enforce-eager` | Disable CUDA graph execution. | +| `--max-cudagraph-capture-size` | Largest batch size to capture with CUDA graphs. | +| `--cudagraph-capture-sizes` | Explicit CUDA graph capture sizes. | + +`--chunked-prefill-size` is intentionally separate from +`--max-num-batched-tokens`: in TokenSpeed it is the scheduler's per-iteration +issue budget, while `--max-total-tokens` controls the global token pool. + +## Parallelism + +| Parameter | Purpose | +| --- | --- | +| `--tensor-parallel-size`, `--tp` | Familiar alias for setting attention tensor parallel size. | +| `--attn-tp-size` | Tensor parallel size for attention. | +| `--dense-tp-size` | Tensor parallel size for dense layers. | +| `--moe-tp-size` | Tensor parallel size for MoE layers. | +| `--data-parallel-size` | Number of data-parallel replicas. | +| `--enable-expert-parallel` | Set expert parallelism across the selected world size. | +| `--expert-parallel-size`, `--ep-size` | Explicit expert parallel size. | +| `--world-size` | Total worker process count across all nodes. | +| `--nprocs-per-node` | Worker process count per node. | +| `--nnodes` | Number of nodes. | +| `--node-rank` | Rank of the current node. | +| `--dist-init-addr` | Distributed initialization address. | + +Use `--tensor-parallel-size` for simple launches. Use the +TokenSpeed-specific split knobs when attention, dense, and MoE layers need +different process groups. + +## Backend Selection + +| Parameter | Purpose | +| --- | --- | +| `--attention-backend` | Attention kernel backend. Common values include `mha`, `fa3`, `fa4`, `triton`, `flashinfer`, `trtllm_mla`, and `tokenspeed_mla`. | +| `--drafter-attention-backend` | Attention backend for speculative decoding drafter model. | +| `--moe-backend` | MoE backend. | +| `--draft-moe-backend` | MoE backend for the speculative decoding draft model. | +| `--all2all-backend` | MoE all-to-all backend. | +| `--deepep-mode` | DeepEP mode: `auto`, `normal`, or `low_latency`. | +| `--sampling-backend` | Sampling backend: `greedy`, `flashinfer`, `flashinfer_full`, `triton`, or `triton_full`. | + +Set backend choices explicitly in production. `auto` is useful for bring-up, but +explicit values make benchmark comparisons and regressions easier to reason +about. + +When `--dp-sampling` is enabled, the logits processor owns the per-forward +logits layout decision and carries the resulting plan to the sampling backend +with the logits output. + +## Reasoning And Tool Calling + +| Parameter | Purpose | +| --- | --- | +| `--reasoning-parser` | Parser for extracting reasoning content from model outputs (handled by the smg gateway). | +| `--tool-call-parser` | Parser for OpenAI-compatible tool-call payloads (handled by the smg gateway). | +| `--enable-custom-logit-processor` | Allow custom logit processors. Keep disabled unless the deployment needs it. | + +Common reasoning parser values include `kimi_k25`, `base`, `qwen3`, +`deepseek_r1`, and `deepseek_v31`. Common tool-call parser values include +`kimik2`, `qwen`, `deepseek_v4`, `json`, and `passthrough`. The parser names +are validated by the SMG gateway, so use +the values accepted by the bundled `tokenspeed-smg` package. + +## Speculative Decoding + +| Parameter | Purpose | +| --- | --- | +| `--speculative-config` | JSON speculative decoding configuration. | +| `--speculative-algorithm` | Speculative algorithm, such as `EAGLE3`, `MTP`, or `DFLASH`. | +| `--speculative-draft-model-path` | Draft model path or repo ID. | +| `--speculative-draft-model-quantization` | Draft model quantization. Defaults to `unquant`. | +| `--speculative-num-steps` | Number of draft model steps. Defaults to `3`. | +| `--speculative-num-draft-tokens` | Number of draft tokens. Defaults to `--speculative-num-steps + 1`. | +| `--speculative-eagle-topk` | EAGLE top-k. Defaults to `1`. | +| `--eagle3-layers-to-capture` | EAGLE3 layers to capture. | + +Prefer `--speculative-config` for recipe-style launches because it keeps method, +draft model, and token count together. + +## Observability + +| Parameter | Purpose | +| --- | --- | +| `--log-level` | Runtime log level. | +| `--log-level-http` | HTTP server log level. Defaults to `--log-level` when unset. | +| `--enable-log-requests` | Log request metadata and optionally payloads. | +| `--log-requests-level` | Request logging verbosity. | +| `--enable-log-request-stats` | Log a one-line per-request performance summary on finish/abort (see below). | +| `--enable-metrics` | Enable metrics reporting. | +| `--metrics-reporters` | Metrics reporter, such as `prometheus`. | +| `--decode-log-interval` | Decode batch log interval. | +| `--enable-cache-report` | Include cached-token counts in OpenAI-compatible usage details. | +| `--kv-events-config` | JSON config for KV cache mutation events. Set `enable_kv_cache_events` and a publisher such as `zmq` to publish device prefix-cache stores and removals. | + +### Per-Request Stats + +`--enable-log-request-stats` enriches the scheduler's per-request finish line for +latency/throughput debugging. When set, the `Req: Finish! ...` line carries +a Python-object repr (`RequestStats(...)`) instead of the default +`Accept_num_tokens_avg` value (which it subsumes as `acc_len`). Every field is +derived from host-side timestamps and counters already available in the +scheduler — it adds **no GPU sync** and so no engine slowdown. Example: + +``` +Req: chatcmpl-019ef6b7 Finish! RequestStats(status='finished', reason='stop', prompt_tokens=28684, cache_tokens=832, output_tokens=33, cache_hit_rate=0.029, queue_ms=13.8, prefill_ms=15.8, ttft_ms=42.1, total_ms=58.0, preempt_ms=0.0, preempt_count=0, decode_tps=210.4, acc_len=None, acc_rate=None, recv_ts=1782255696.726, commit_ts=1782255696.74, finish_ts=1782255696.784) +``` + +| Field | Meaning | +| --- | --- | +| `status` / `reason` | `finished` vs `aborted`; finish-reason type (`stop`/`length`/`abort`). | +| `prompt_tokens` / `cache_tokens` / `output_tokens` | Prompt tokens, prefix-cache-hit tokens, generated tokens. | +| `cache_hit_rate` | `cache_tokens / prompt_tokens` (0–1). | +| `queue_ms` | Received → first scheduled into a forward batch. | +| `prefill_ms` | Scheduled → prefill complete. | +| `ttft_ms` | Received → first output token (always ≥ `prefill_ms`; it also spans the queue). | +| `total_ms` | Received → finished/aborted. | +| `preempt_ms` / `preempt_count` | Wall-clock this request's decode was delayed by prefilling other requests, and the number of such interruptions. Host-side best-effort. | +| `decode_tps` | Decode throughput (generated tokens / decode window). | +| `acc_len` / `acc_rate` | Spec-decode acceptance length and rate (`None` when speculative decoding is off). | +| `recv_ts` / `commit_ts` / `finish_ts` | Absolute epoch timestamps for received / scheduled / finished. | + +### KV Cache Events + +KV cache events publish reusable device prefix-cache mutations from the live +C++ scheduler path. Host/L2 loadback events are not published by this initial +stream. Block hash lineage is cached on prefix-cache nodes, so publishing a +stored block uses the parent node's cached hash instead of rebuilding the full +ancestor prefix. + +Example: + +```bash +--kv-events-config '{"enable_kv_cache_events":true,"publisher":"zmq","endpoint":"tcp://*:5557","topic":"kv-events"}' +``` + +The ZMQ publisher sends three frames: topic bytes, an 8-byte big-endian sequence +number, and a msgpack payload. The payload is an array-like `KVEventBatch`: + +```python +[timestamp, [["BlockStored", [block_hash], parent_hash, token_ids, block_size]], attn_dp_rank] +[timestamp, [["BlockRemoved", [block_hash]]], attn_dp_rank] +``` + +With attention data parallelism, each attention DP rank publishes on an offset +port from the configured endpoint. + +## TokenSpeed-Specific Runtime Knobs + +These parameters are TokenSpeed-specific. They expose runtime +features directly: + +- `--max-total-tokens` +- `--max-prefill-tokens` +- `--chunked-prefill-size` +- `--attn-tp-size` +- `--dense-tp-size` +- `--moe-tp-size` +- `--kvstore-*` +- `--enable-mla-l1-5-cache` +- `--kv-events-config` +- `--mla-chunk-multiplier` +- `--disaggregation-*` +- `--comm-fusion-max-num-tokens` +- `--enable-allreduce-fusion` diff --git a/docs/guides/getting-started.md b/docs/guides/getting-started.md new file mode 100644 index 0000000..7729ffb --- /dev/null +++ b/docs/guides/getting-started.md @@ -0,0 +1,76 @@ +# Getting Started + +This guide brings up a TokenSpeed development environment and verifies that the +runtime can start. + +## Prerequisites + +- NVIDIA GPU host +- Docker with GPU support +- enough shared memory for model serving +- access to the model checkpoints you plan to serve + +## Start a Runner Container + +```bash +docker pull lightseekorg/tokenspeed-runner:latest + +docker run -itd \ + --shm-size 32g \ + --gpus all \ + -v /raid/cache:/home/runner/.cache \ + --ipc=host \ + --network=host \ + --pid=host \ + --privileged \ + --name tokenspeed \ + lightseekorg/tokenspeed-runner:latest \ + /bin/bash +``` + +Inside the container: + +```bash +git clone https://github.com/lightseekorg/tokenspeed.git +cd tokenspeed +``` + +## Install Packages + +Install the Python runtime: + +```bash +export PIP_BREAK_SYSTEM_PACKAGES=1 +pip install -e "./python" --no-build-isolation +``` + +Install the kernel package. Its Python package metadata installs the selected +backend dependencies automatically. + +```bash +pip install -e tokenspeed-kernel/python/ --no-build-isolation +``` + +Install the scheduler package: + +```bash +pip install -e tokenspeed-scheduler/ +``` + +## Verify + +```bash +tokenspeed env +tokenspeed serve --help +``` + +## Launch + +```bash +tokenspeed serve openai/gpt-oss-20b \ + --host 0.0.0.0 \ + --port 8000 \ + --tensor-parallel-size 1 +``` + +For model-specific examples, continue with [Model Recipes](../recipes/models.md). diff --git a/docs/guides/launching.md b/docs/guides/launching.md new file mode 100644 index 0000000..6a2d76b --- /dev/null +++ b/docs/guides/launching.md @@ -0,0 +1,65 @@ +# Launching a Server + +`tokenspeed serve` starts an OpenAI-compatible HTTP server. Put the model path +directly after the command. + +## Minimal Launch + +```bash +tokenspeed serve openai/gpt-oss-20b \ + --host 0.0.0.0 \ + --port 8000 \ + --tensor-parallel-size 1 +``` + +## Production Launch Skeleton + +Use explicit parameters in scripts so a deployment is reproducible. + +```bash +tokenspeed serve nvidia/Kimi-K2.5-NVFP4 \ + --served-model-name kimi-k2.5 \ + --host 0.0.0.0 \ + --port 8000 \ + --trust-remote-code \ + --max-model-len 262144 \ + --kv-cache-dtype fp8 \ + --quantization nvfp4 \ + --tensor-parallel-size 4 \ + --enable-expert-parallel \ + --chunked-prefill-size 8192 \ + --max-num-seqs 256 \ + --attention-backend trtllm_mla \ + --moe-backend flashinfer_trtllm \ + --reasoning-parser kimi_k25 \ + --tool-call-parser kimik2 +``` + +## Launch Checklist + +- Put the model path directly after `tokenspeed serve`. +- Set `--host`, `--port`, and `--served-model-name` for the API surface. +- Set `--max-model-len`, `--kv-cache-dtype`, and `--gpu-memory-utilization` before tuning concurrency. +- Set `--tensor-parallel-size`, `--data-parallel-size`, and `--enable-expert-parallel` to match the hardware topology. +- Set model-family parsers such as `--reasoning-parser` and `--tool-call-parser` when the chat format needs them. +- Set backend choices explicitly for benchmark or production runs. + +## OpenAI-Compatible Client + +```python +from openai import OpenAI + +client = OpenAI(api_key="EMPTY", base_url="http://localhost:8000/v1") +response = client.chat.completions.create( + model="kimi-k2.5", + messages=[{"role": "user", "content": "Write a concise deployment checklist."}], + max_tokens=256, +) +print(response.choices[0].message.content) +``` + +## Next + +- [Model Recipes](../recipes/models.md) +- [Server Parameters](../configuration/server.md) +- [Parallelism](../serving/parallelism.md) diff --git a/docs/index.md b/docs/index.md new file mode 100644 index 0000000..b41fef0 --- /dev/null +++ b/docs/index.md @@ -0,0 +1,69 @@ +--- +layout: home + +hero: + name: TokenSpeed + text: Speed-of-light LLM inference + tagline: Production-oriented docs for launching, tuning, and operating low-latency OpenAI-compatible serving. + actions: + - theme: brand + text: Get Started + link: /guides/getting-started + - theme: alt + text: Launch Recipes + link: /recipes/models + - theme: alt + text: Server Parameters + link: /configuration/server + +features: + - title: Launch First + details: Start with concrete commands, then tune the exact knobs that affect memory, scheduling, parallelism, and kernels. + - title: Familiar Parameters + details: TokenSpeed keeps familiar parameter names where the runtime semantics match, with TokenSpeed-specific knobs documented separately. + - title: Model Recipes + details: Recipes collect the launch patterns used for Kimi and GPT-OSS deployments. + - title: Operational Surface + details: Parallelism and configuration guidance stay close to the serving paths operators actually use. +--- + +## Start Here + +- [Getting Started](./guides/getting-started.md) +- [Launching a Server](./guides/launching.md) +- [Model Recipes](./recipes/models.md) +- [Server Parameters](./configuration/server.md) +- [Compatible Parameters](./configuration/compatible-parameters.md) +- [Parallelism](./serving/parallelism.md) + +## Common Workflow + +1. Install the runtime and kernel packages. +2. Pick a launch recipe close to your model family and hardware. +3. Set model loading, memory, scheduler, and parallelism parameters explicitly. +4. Validate correctness and throughput together before changing more than one tuning dimension. + +## Minimal Server + +```bash +tokenspeed serve openai/gpt-oss-20b \ + --host 0.0.0.0 \ + --port 8000 \ + --tensor-parallel-size 1 +``` + +The server exposes an OpenAI-compatible API under `/v1`. + +## High-Performance Shape + +Large MoE deployments usually make the same decisions: + +- model path and revision +- context length and KV cache dtype +- scheduler token and sequence budgets +- attention and MoE backends +- tensor, data, and expert parallelism +- reasoning, tool-call, and speculative decoding parsers + +See [Model Recipes](./recipes/models.md) for concrete examples and +[Server Parameters](./configuration/server.md) for the parameter reference. diff --git a/docs/package.json b/docs/package.json new file mode 100644 index 0000000..528e074 --- /dev/null +++ b/docs/package.json @@ -0,0 +1,13 @@ +{ + "name": "tokenspeed-docs", + "private": true, + "packageManager": "bun@1.3.0", + "scripts": { + "docs:dev": "vitepress dev", + "docs:build": "vitepress build", + "docs:preview": "vitepress preview" + }, + "devDependencies": { + "vitepress": "^1.6.3" + } +} diff --git a/docs/recipes/models.md b/docs/recipes/models.md new file mode 100644 index 0000000..af7b858 --- /dev/null +++ b/docs/recipes/models.md @@ -0,0 +1,213 @@ +# Model Recipes + +These recipes start from a known model family, pick the hardware topology, then +set only the parameters that change runtime behavior. + +The commands below are templates. Validate exact model IDs, checkpoint formats, +and backend choices against the build you deploy. + +## Kimi K2.5 / K2.6 + +Kimi-style MoE launches usually need remote code, long context, reasoning and +tool parsers, and explicit MLA/MoE backends. + +```bash +tokenspeed serve nvidia/Kimi-K2.5-NVFP4 \ + --served-model-name kimi-k2.5 \ + --trust-remote-code \ + --max-model-len 262144 \ + --kv-cache-dtype fp8 \ + --quantization nvfp4 \ + --tensor-parallel-size 4 \ + --enable-expert-parallel \ + --chunked-prefill-size 8192 \ + --max-num-seqs 256 \ + --attention-backend trtllm_mla \ + --moe-backend flashinfer_trtllm \ + --reasoning-parser kimi_k25 \ + --tool-call-parser kimik2 \ + --host 0.0.0.0 \ + --port 8000 +``` + +For K2.6, keep the same parameter shape and change the checkpoint and parser +only if the model card requires a different value. + +To enable a compatible DFlash draft model, keep the target launch shape and add +the draft model path plus DFlash speculative decoding options: + +```bash +tokenspeed serve nvidia/Kimi-K2.6-NVFP4 \ + --served-model-name kimi-k2.6 \ + --trust-remote-code \ + --max-model-len 262144 \ + --kv-cache-dtype fp8 \ + --quantization nvfp4 \ + --tensor-parallel-size 4 \ + --enable-expert-parallel \ + --chunked-prefill-size 8192 \ + --max-num-seqs 256 \ + --attention-backend tokenspeed_mla \ + --moe-backend flashinfer_trtllm \ + --reasoning-parser kimi_k25 \ + --tool-call-parser kimik2 \ + --speculative-algorithm DFLASH \ + --speculative-draft-model-path /path/to/kimi-k2.6-dflash \ + --speculative-num-draft-tokens 8 \ + --speculative-num-steps 7 \ + --drafter-attention-backend fa4 \ + --host 0.0.0.0 \ + --port 8000 +``` + +Known limitation: native TokenSpeed DFlash currently uses full-history draft +attention. It does not yet expose an equivalent of SGLang's +`--speculative-dflash-draft-window-size`; add such a flag before relying on +bounded draft attention for long-context deployments. + +## GLM5 / GLM5.2 + +GLM5 launches usually need remote code, long context, expert parallelism, FP8 KV +cache, and the TRTLLM MoE backend. GLM5.2 FP8 is available on Hugging Face as +`zai-org/GLM-5.2-FP8`. TokenSpeed defaults the reasoning parser to `glm45`; +pass an explicit parser flag to override it. + +```bash +tokenspeed serve zai-org/GLM-5.2-FP8 \ + --served-model-name glm-5.2 \ + --trust-remote-code \ + --tensor-parallel-size 8 \ + --enable-expert-parallel \ + --moe-backend flashinfer_trtllm \ + --kv-cache-dtype fp8 \ + --max-model-len 262144 \ + --chunked-prefill-size 8192 \ + --max-num-seqs 128 \ + --host 0.0.0.0 \ + --port 8000 +``` + +## Qwen3 Dense / Qwen3 30B-A3B + +Qwen2, dense Qwen3, and Qwen3 MoE checkpoints use different architecture names. +For Qwen3 30B-A3B, the Hugging Face config advertises `qwen3_moe` and +`Qwen3MoeForCausalLM`, so launch it as a MoE model. + +```bash +tokenspeed serve Qwen/Qwen3-30B-A3B \ + --served-model-name qwen3-30b-a3b \ + --tensor-parallel-size 2 \ + --enable-expert-parallel \ + --moe-backend flashinfer_cutlass \ + --max-model-len 40960 \ + --reasoning-parser qwen3 \ + --host 0.0.0.0 \ + --port 8000 +``` + +## GPT-OSS 20B / 120B + +Small GPT-OSS launches can start simple. Large GPT-OSS launches usually tune +tensor parallelism, scheduler token budget, and KV cache dtype. + +```bash +tokenspeed serve openai/gpt-oss-20b \ + --served-model-name gpt-oss-20b \ + --tensor-parallel-size 1 \ + --max-model-len 131072 \ + --chunked-prefill-size 8192 \ + --reasoning-parser base \ + --host 0.0.0.0 \ + --port 8000 +``` + +```bash +tokenspeed serve openai/gpt-oss-120b \ + --served-model-name gpt-oss-120b \ + --tensor-parallel-size 4 \ + --max-model-len 131072 \ + --kv-cache-dtype fp8 \ + --chunked-prefill-size 8192 \ + --max-num-seqs 256 \ + --reasoning-parser base \ + --host 0.0.0.0 \ + --port 8000 +``` + +## DeepSeek V4-Flash / V4-Pro + +DeepSeek V4 needs FP8 KV cache, the DeepGEMM `mega_moe` experts, and the FP4 +indexer cache. `tokenspeed serve` auto-selects `--reasoning-parser deepseek_v31` +and `--tool-call-parser deepseek_v4`, and auto-sets `block_size=256` (pass +`--block-size N` with `N != 64` to override). Requires +`tokenspeed-deepgemm>=2.5.0.post20260629` and `tokenspeed-flashmla`. + +**V4-Flash** — 4× B200 (SM100), data-parallel + expert-parallel: + +```bash +tokenspeed serve deepseek-ai/DeepSeek-V4-Flash \ + --served-model-name deepseek-v4-flash \ + --trust-remote-code \ + --data-parallel-size 4 \ + --enable-expert-parallel \ + --kv-cache-dtype fp8_e4m3 \ + --moe-backend mega_moe \ + --attention-use-fp4-indexer-cache \ + --max-model-len 80000 \ + --max-total-tokens 163840 \ + --chunked-prefill-size 8192 \ + --enable-mixed-batch \ + --gpu-memory-utilization 0.9 \ + --disable-kvstore \ + --host 0.0.0.0 \ + --port 8000 +``` + +**V4-Pro** — 8× B200, tensor-parallel: + +```bash +tokenspeed serve deepseek-ai/DeepSeek-V4-Pro \ + --served-model-name deepseek-v4-pro \ + --trust-remote-code \ + --tensor-parallel-size 8 \ + --kv-cache-dtype fp8_e4m3 \ + --moe-backend flashinfer_trtllm \ + --attention-use-fp4-indexer-cache \ + --max-model-len 80000 \ + --max-total-tokens 2560000 \ + --chunked-prefill-size 8192 \ + --gpu-memory-utilization 0.9 \ + --disable-kvstore \ + --host 0.0.0.0 \ + --port 8000 +``` + +For the expert-parallel topology, swap `--tensor-parallel-size 8` for +`--tensor-parallel-size 8 --enable-expert-parallel --dense-tp-size 1` and +`--moe-backend flashinfer_trtllm` for `--moe-backend mega_moe`. + +### MTP speculative decoding + +Both variants can drive the checkpoint's NextN/MTP draft layers. Keep the launch +flags above and add: + +```bash +--speculative-algorithm MTP \ +--speculative-num-steps 3 +``` + +With `--speculative-draft-model-path` omitted, V4 uses the same checkpoint as the +draft source (`DeepseekV4ForCausalLMNextN`). MTP runs on the non-overlap +scheduler — the runtime disables overlap scheduling automatically when +speculative decoding and paged-cache groups are both active — and prefix caching +stays on by default. Add `--enable-metrics` to read `Decoded Tok/Iter` and the +speculative accept rate from the run summary. + +## Tuning Order + +1. Set model ID, trust policy, tokenizer mode, and served model name. +2. Set context length and KV cache dtype. +3. Set tensor, data, and expert parallelism to match the node topology. +4. Set scheduler budgets: `--chunked-prefill-size`, `--max-num-seqs`, and only then `--max-total-tokens`. +5. Set attention, MoE, and sampling backends explicitly for benchmark runs. +6. Add reasoning, tool-call, grammar, or speculative decoding only when the model and workload need them. diff --git a/docs/serving/parallelism.md b/docs/serving/parallelism.md new file mode 100644 index 0000000..884932f --- /dev/null +++ b/docs/serving/parallelism.md @@ -0,0 +1,81 @@ +# Parallelism + +TokenSpeed exposes familiar `--tensor-parallel-size` and `--tp` entry points +plus additional split parallelism controls for attention, dense, and MoE layers. + +## Quick Start + +Use this form when the same tensor-parallel group is acceptable for the model: + +```bash +tokenspeed serve \ + --tensor-parallel-size 8 +``` + +`--tensor-parallel-size` maps to TokenSpeed attention tensor parallelism and +cannot be used together with `--attn-tp-size`. + +## Split Parallelism + +Use split knobs when different layer families need different process groups: + +```bash +tokenspeed serve \ + --world-size 8 \ + --attn-tp-size 4 \ + --dense-tp-size 4 \ + --moe-tp-size 4 +``` + +| Parameter | Use | +| --- | --- | +| `--world-size` | Total worker processes across all nodes. | +| `--nprocs-per-node` | Worker processes launched on each node. | +| `--attn-tp-size` | Attention tensor parallel size. | +| `--dense-tp-size` | Dense layer tensor parallel size. | +| `--moe-tp-size` | MoE layer tensor parallel size. | +| `--data-parallel-size` | Replicated data-parallel groups. | +| `--enable-expert-parallel` | Expert parallelism across the selected world size. | +| `--expert-parallel-size` | Explicit expert parallel size. | + +## MoE Deployments + +Large MoE models usually choose one of these shapes: + +- TP only: simplest startup path, often best for smaller MoE checkpoints. +- TP + EP: tensor parallelism within a replica, expert parallelism across ranks. +- DP + EP: multiple replicated decode groups with experts distributed inside each group. + +Start with the recipe closest to your model family, then tune: + +- `--tensor-parallel-size` or split TP values +- `--enable-expert-parallel` +- `--moe-backend` +- `--all2all-backend` +- `--deepep-mode` + +## Multi-Node + +Set these explicitly: + +```bash +tokenspeed serve \ + --nnodes 2 \ + --node-rank 0 \ + --nprocs-per-node 8 \ + --world-size 16 \ + --dist-init-addr :25000 +``` + +Each node must use the same model, backend, precision, and scheduler settings. +Only `--node-rank` should differ between nodes. + +## Validation + +Before benchmarking: + +- verify every rank starts and joins the distributed group +- verify the API responds before sending load +- confirm GPU visibility and process placement +- compare output correctness before tuning throughput +- keep the full launch command with benchmark results diff --git a/python/LICENSE b/python/LICENSE new file mode 100644 index 0000000..e95f89c --- /dev/null +++ b/python/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 LightSeek Foundation + +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. diff --git a/python/THIRDPARTYNOTICES b/python/THIRDPARTYNOTICES new file mode 100644 index 0000000..07d81ac --- /dev/null +++ b/python/THIRDPARTYNOTICES @@ -0,0 +1,271 @@ +Portions of this project are based on or inspired by designs and code from the following projects: + +Notice for NVIDIA/TensorRT-LLM +------------------------------- +Apache License, Version 2.0 + +Copyright contributors to the TensorRT LLM project + +Licensed under the Apache License, Version 2.0, whose full license text is available below. + +Notice for FluentLLM +------------------------------- +Apache License, Version 2.0 + +Copyright contributors to the FluentLLM project + +Licensed under the Apache License, Version 2.0, whose full license text is available below. + +Notice for vllm-project/vllm +------------------------------- +Apache License, Version 2.0 + +Copyright contributors to the vLLM project + +Licensed under the Apache License, Version 2.0, whose full license text is available below. + +Notice for SGLang +------------------------------- +Apache License, Version 2.0 + +Copyright 2023-2024 SGLang Team + +Licensed under the Apache License, Version 2.0, whose full license text is available below. + +Notice for Qwen +------------------------------- +Apache License, Version 2.0 + +Copyright 2025 The Qwen team, Alibaba Group + +Licensed under the Apache License, Version 2.0, whose full license text is available below. + +Notice for deepseek-ai/EPLB +------------------------------- +MIT License + +Copyright (c) 2025 DeepSeek + +Licensed under the MIT License, whose full license text is available in the repository LICENSE file. + +Notice for HuggingFace Transformers +------------------------------- +Apache License, Version 2.0 + +Copyright 2025 HuggingFace Inc. team + +Licensed under the Apache License, Version 2.0, whose full license text is available below. + +Notice for fla-org/flash-linear-attention +------------------------------- +MIT License + +Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li + +Licensed under the MIT License, whose full license text is available in the repository LICENSE file. + +================================================================================ + Apache 2.0 LICENSE +================================================================================ + + 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 [yyyy] [name of copyright owner] + + 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/python/pyproject.toml b/python/pyproject.toml new file mode 100755 index 0000000..689672e --- /dev/null +++ b/python/pyproject.toml @@ -0,0 +1,97 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +[build-system] +requires = ["setuptools==69.5.1", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "tokenspeed" +version = "0.1.0" +description = "TokenSpeed is a speed-of-light LLM inference engine." +readme = "README.md" +requires-python = ">=3.10" +license = { file = "LICENSE" } +classifiers = [ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: MIT License", +] +dependencies = [ + "aiohttp", + "compressed-tensors", + "dill", + "einops", + "fastapi", + "hf_transfer", + "huggingface_hub", + "modelscope", + "msgspec", + "ninja", + "numpy", + "openai>=2.24.1", + "openai-harmony", + "orjson", + "packaging", + "partial-json-parser", + "peft", + "pillow", + "prometheus-client", + "psutil", + "pybase64", + "pybind11", + "pydantic", + "py-spy", + "pytest-asyncio", + "python-multipart", + "pyzmq", + "requests", + "setproctitle", + "tiktoken", + "tokenspeed-mooncake>=0.3.11.post20260527", + "tokenspeed-smg==1.7.0.post20260710", + "tokenspeed-smg-grpc-proto==0.4.12.post20260710", + "tokenspeed-smg-grpc-servicer==0.6.0.post20260710", + "torch==2.11.0", + # Tag-capable sleep/wake allocator (region(tag, enable_cpu_backup)/pause/ + # resume). Bundles cu12/cu13 binaries; import-guarded so it's a no-op unless + # --enable-memory-saver is set. See runtime/utils/torch_memory_saver_adapter.py. + "torch_memory_saver==0.0.9.post1", + "torchvision", + "tqdm", + "transformers==5.12.0", + "uv", + "uvicorn", + "uvloop", + "xgrammar==0.2.2", + "viztracer", +] + +[project.scripts] +tokenspeed = "tokenspeed.cli:main" +ts = "tokenspeed.cli:main" + +[project.urls] +"Homepage" = "https://github.com/lightseekorg/tokenspeed" + +[tool.setuptools] +license-files = ["LICENSE", "THIRDPARTYNOTICES"] + +[tool.setuptools.packages.find] +include = ["tokenspeed*"] diff --git a/python/tokenspeed/__init__.py b/python/tokenspeed/__init__.py new file mode 100755 index 0000000..a32d61e --- /dev/null +++ b/python/tokenspeed/__init__.py @@ -0,0 +1,25 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from tokenspeed._logging import suppress_noisy_third_party_logs + +_suppress_noisy_third_party_logs = suppress_noisy_third_party_logs + +_suppress_noisy_third_party_logs() diff --git a/python/tokenspeed/_logging.py b/python/tokenspeed/_logging.py new file mode 100644 index 0000000..032942e --- /dev/null +++ b/python/tokenspeed/_logging.py @@ -0,0 +1,110 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import logging +import os +import warnings +from importlib import import_module + +_original_showwarning = warnings.showwarning + +_CUTLASS_POINTER_WARNING = "Use explicit `struct.scalar.ptr` for pointer instead." +_CUTLASS_NAMED_BARRIER_WARNING = ( + "NamedBarrier wait also arrives on the barrier. " + "Routing call to NamedBarrier.arrive_and_wait()." +) + + +def _is_noisy_cutlass_dsl_warning(message, category) -> bool: + message_text = str(message) + return ( + issubclass(category, DeprecationWarning) + and message_text == _CUTLASS_POINTER_WARNING + ) or ( + issubclass(category, UserWarning) + and message_text == _CUTLASS_NAMED_BARRIER_WARNING + ) + + +def _showwarning(message, category, filename, lineno, file=None, line=None): + if _is_noisy_cutlass_dsl_warning(message, category): + return + _original_showwarning(message, category, filename, lineno, file=file, line=line) + + +def _suppress_cutlass_dsl_warnings(): + if warnings.showwarning is not _showwarning: + warnings.showwarning = _showwarning + + warnings.filterwarnings( + "ignore", + message=r"Use explicit `struct\.scalar\.ptr` for pointer instead\.", + category=DeprecationWarning, + ) + warnings.filterwarnings( + "ignore", + message=( + r"NamedBarrier wait also arrives on the barrier\. " + r"Routing call to NamedBarrier\.arrive_and_wait\(\)\." + ), + category=UserWarning, + ) + + +def _suppress_flash_attn_jit_cache_debug_log(): + logger_name = "flash_attn.cute.cache_utils" + previous_disable_level = logging.root.manager.disable + logging.disable(logging.INFO) + try: + import_module(logger_name) + except Exception: + return + finally: + logging.disable(previous_disable_level) + + logger = logging.getLogger(logger_name) + logger.setLevel(logging.WARNING) + for handler in logger.handlers: + handler.setLevel(logging.WARNING) + + +def suppress_noisy_third_party_logs(): + os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1") + os.environ.setdefault("TLLM_LOG_LEVEL", "WARNING") + _suppress_cutlass_dsl_warnings() + + for logger_name in ( + "transformers", + "huggingface_hub", + "huggingface_hub.file_download", + "httpx", + "httpcore", + "flash_attn.cute.cache_utils", + ): + logging.getLogger(logger_name).setLevel(logging.WARNING) + + _suppress_flash_attn_jit_cache_debug_log() + + try: + from huggingface_hub.utils import disable_progress_bars + + disable_progress_bars() + except Exception: + pass diff --git a/python/tokenspeed/bench.py b/python/tokenspeed/bench.py new file mode 100755 index 0000000..c9a61f3 --- /dev/null +++ b/python/tokenspeed/bench.py @@ -0,0 +1,2060 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +r"""Benchmark online serving throughput. + +On the server side, launch a TokenSpeed OpenAI-compatible API server: + tokenspeed serve --model + +On the client side, run: + tokenspeed bench serve \ + --backend \ + --label \ + --model \ + --dataset-name \ + --input-len \ + --output-len \ + --request-rate \ + --num-prompts +""" + +from __future__ import annotations + +import argparse +import asyncio +import codecs +import contextlib +import json +import logging +import math +import os +import random +import resource +import ssl +import sys +import time +import traceback +import warnings +from collections.abc import AsyncGenerator, Coroutine +from dataclasses import dataclass, field +from datetime import datetime +from enum import Enum +from typing import Any, Literal +from urllib.parse import urlparse + +import aiohttp +import numpy as np +import requests +from tqdm.asyncio import tqdm +from transformers import AutoTokenizer, PreTrainedTokenizerBase + +from tokenspeed.runtime.utils.env import envs + +# Streaming HTTP timeouts. ``total=6h`` keeps the session umbrella generous so +# whole-run benches don't get cut off; the per-socket sub-timeouts catch a +# legitimately stuck stream without false-failing slow legitimate prefills. +# +# ``sock_read`` defaults to 30 minutes — well above the largest TTFT one would +# expect on real hardware (a 64k-context prefill on a single-GPU consumer card +# is still well under 10 minutes) yet far below ``total``, so an indefinitely +# silent socket still surfaces as a ``aiohttp.ServerTimeoutError`` rather than +# blocking the outer ``asyncio.gather`` at high concurrency. Long-haul or +# pathologically large prefill workloads can bump it via env. ``sock_connect`` +# is the dial-tone timeout for the TCP handshake itself. +AIOHTTP_TOTAL_TIMEOUT_SEC = float( + os.environ.get("TOKENSPEED_BENCH_TOTAL_TIMEOUT_SEC", str(6 * 60 * 60)) +) +AIOHTTP_SOCK_CONNECT_TIMEOUT_SEC = float( + os.environ.get("TOKENSPEED_BENCH_SOCK_CONNECT_TIMEOUT_SEC", "30") +) +AIOHTTP_SOCK_READ_TIMEOUT_SEC = float( + os.environ.get("TOKENSPEED_BENCH_SOCK_READ_TIMEOUT_SEC", str(30 * 60)) +) +AIOHTTP_TIMEOUT = aiohttp.ClientTimeout( + total=AIOHTTP_TOTAL_TIMEOUT_SEC, + sock_connect=AIOHTTP_SOCK_CONNECT_TIMEOUT_SEC, + sock_read=AIOHTTP_SOCK_READ_TIMEOUT_SEC, +) +# Per-request hard ceiling so a single misbehaving stream cannot block the +# whole gather. 1h is generous enough for the longest practical decode and +# still bounded for CI / smoke benches. Override via env when running +# unusually long sequences. +PER_REQUEST_TIMEOUT_SEC = float( + os.environ.get("TOKENSPEED_BENCH_PER_REQUEST_TIMEOUT_SEC", str(60 * 60)) +) +DEFAULT_NUM_PROMPTS = 1000 +MILLISECONDS_TO_SECONDS_CONVERSION = 1000 +SHAREGPT_URL = "https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json" +OPENAI_COMPATIBLE_BACKENDS = frozenset({"openai", "tokenspeed"}) +logger = logging.getLogger(__name__) + +# Type alias: a single float applies to both ISL and OSL; a dict allows +# specifying them independently via ``{"input": ..., "output": ...}``. +RangeRatio = float | dict[str, float] + + +def _print_section_header(title: str, fill: str) -> None: + print(f"{title:{fill}^50}") + + +def _print_metric_row(label: str, value: Any, precision: int | None = None) -> None: + formatted_value = ( + f"{value:<10}" if precision is None else f"{value:<10.{precision}f}" + ) + print(f"{label:<40} {formatted_value}") + + +class StreamedResponseHandler: + """Accumulate SSE bytes until complete `data:` messages are available.""" + + def __init__(self) -> None: + self.buffer = "" + self._decoder = codecs.getincrementaldecoder("utf-8")() + + def add_chunk(self, chunk_bytes: bytes) -> list[str]: + self.buffer += self._decoder.decode(chunk_bytes) + messages: list[str] = [] + + while "\n\n" in self.buffer: + message, self.buffer = self.buffer.split("\n\n", 1) + message = message.strip() + if message: + messages.append(message) + + if self.buffer.startswith("data: "): + message_content = self.buffer.removeprefix("data: ").strip() + if message_content == "[DONE]": + messages.append(self.buffer.strip()) + self.buffer = "" + elif message_content: + try: + json.loads(message_content) + except json.JSONDecodeError: + pass + else: + messages.append(self.buffer.strip()) + self.buffer = "" + + return messages + + +@dataclass +class SampleRequest: + prompt: str + prompt_len: int + expected_output_len: int + multi_modal_data: dict | list[dict] | None = None + lora_request: Any | None = None + request_id: str | None = None + + +@dataclass +class RequestFuncInput: + """The input for the request function.""" + + prompt: str | list[str] + api_url: str + prompt_len: int + output_len: int + model: str + model_name: str | None = None + logprobs: int | None = None + extra_headers: dict | None = None + extra_body: dict | None = None + multi_modal_content: dict | list[dict] | None = None + ignore_eos: bool = False + language: str | None = None + request_id: str | None = None + + +@dataclass +class RequestFuncOutput: + """The output of the request function including metrics.""" + + generated_text: str = "" + success: bool = False + latency: float = 0.0 + output_tokens: int = 0 + ttft: float = 0.0 # Time to first token + itl: list[float] = field(default_factory=list) # list of inter-token latencies + tpot: float = 0.0 # avg next-token latencies + prompt_len: int = 0 + error: str = "" + start_time: float = 0.0 + input_audio_duration: float = 0.0 # in seconds + + +async def await_with_per_request_timeout( + coro: Coroutine[Any, Any, RequestFuncOutput], + *, + prompt_len: int, + pbar: tqdm | None = None, +) -> RequestFuncOutput: + """Run a request coroutine under :data:`PER_REQUEST_TIMEOUT_SEC`. + + Wraps the per-request ``asyncio.wait_for`` so a single stuck stream + cannot deadlock the outer ``asyncio.gather`` in :func:`benchmark`. On + :class:`asyncio.TimeoutError`, returns a standard + :class:`RequestFuncOutput` with ``success=False`` so the gather can + complete and the metrics output reports the failure normally. + """ + try: + return await asyncio.wait_for(coro, timeout=PER_REQUEST_TIMEOUT_SEC) + except asyncio.TimeoutError: + output = RequestFuncOutput() + output.prompt_len = prompt_len + output.success = False + output.error = ( + f"per-request timeout {PER_REQUEST_TIMEOUT_SEC:.1f}s " + "(TOKENSPEED_BENCH_PER_REQUEST_TIMEOUT_SEC)" + ) + if pbar is not None: + pbar.update(1) + return output + + +class TaskType(Enum): + GENERATION = "generation" + + +@dataclass +class BenchmarkMetrics: + completed: int + failed: int + total_input: int + total_output: int + request_throughput: float + request_goodput: float + output_throughput: float + total_token_throughput: float + mean_ttft_ms: float + median_ttft_ms: float + std_ttft_ms: float + percentiles_ttft_ms: list[tuple[float, float]] + mean_tpot_ms: float + median_tpot_ms: float + std_tpot_ms: float + percentiles_tpot_ms: list[tuple[float, float]] + mean_itl_ms: float + median_itl_ms: float + std_itl_ms: float + percentiles_itl_ms: list[tuple[float, float]] + mean_e2el_ms: float + median_e2el_ms: float + std_e2el_ms: float + percentiles_e2el_ms: list[tuple[float, float]] + max_output_tokens_per_s: float + max_concurrent_requests: int + + +def set_ulimit(target_soft_limit: int = 65535) -> None: + 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}") + + +def join_host_port(host: str, port: int) -> str: + return ( + f"[{host}]:{port}" + if ":" in host and not host.startswith("[") + else f"{host}:{port}" + ) + + +def _validate_api_url( + api_url: str, + api_name: str, + expected_suffixes: str | set[str], +) -> None: + if isinstance(expected_suffixes, str): + expected_suffixes = {expected_suffixes} + + expected_suffixes = {*expected_suffixes, "profile"} + + if not api_url.endswith(tuple(expected_suffixes)): + raise ValueError(f"{api_name} URL must end with one of: {expected_suffixes}.") + + +def _update_payload_common( + payload: dict[str, Any], + request_func_input: RequestFuncInput, +) -> None: + if request_func_input.ignore_eos: + payload["ignore_eos"] = request_func_input.ignore_eos + if request_func_input.extra_body: + payload.update(request_func_input.extra_body) + + +def _update_headers_common( + headers: dict[str, Any], + request_func_input: RequestFuncInput, +) -> None: + if request_func_input.extra_headers: + headers |= request_func_input.extra_headers + if request_func_input.request_id: + headers["x-request-id"] = request_func_input.request_id + + +def _get_headers(content_type: str | None = None) -> dict[str, str]: + headers = {} + if content_type: + headers["Content-Type"] = content_type + api_key = os.environ.get("OPENAI_API_KEY") + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + return headers + + +async def async_request_openai_completions( + request_func_input: RequestFuncInput, + session: aiohttp.ClientSession, + pbar: tqdm | None = None, +) -> RequestFuncOutput: + """The async request function for the OpenAI Completions API. + + Args: + request_func_input: The input for the request function. + pbar: The progress bar to display the progress. + + Returns: + The output of the request function. + """ + api_url = request_func_input.api_url + _validate_api_url(api_url, "OpenAI Completions API", "completions") + + payload = { + "model": ( + request_func_input.model_name + if request_func_input.model_name + else request_func_input.model + ), + "prompt": request_func_input.prompt, + "repetition_penalty": 1.0, + "max_tokens": request_func_input.output_len, + "logprobs": request_func_input.logprobs, + "stream": True, + "stream_options": { + "include_usage": True, + }, + } + _update_payload_common(payload, request_func_input) + + headers = _get_headers() + _update_headers_common(headers, request_func_input) + + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + generated_text = "" + st = time.perf_counter() + output.start_time = st + most_recent_timestamp = st + try: + async with session.post(url=api_url, json=payload, headers=headers) as response: + if response.status == 200: + first_chunk_received = False + handler = StreamedResponseHandler() + + async for chunk_bytes in response.content.iter_any(): + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + messages = handler.add_chunk(chunk_bytes) + for message in messages: + if message.startswith(":"): + continue + + chunk = message.removeprefix("data: ") + + if chunk != "[DONE]": + data = json.loads(chunk) + + if choices := data.get("choices"): + text = choices[0].get("text") + timestamp = time.perf_counter() + if not first_chunk_received: + first_chunk_received = True + ttft = time.perf_counter() - st + output.ttft = ttft + else: + output.itl.append(timestamp - most_recent_timestamp) + + most_recent_timestamp = timestamp + generated_text += text or "" + elif usage := data.get("usage"): + output.output_tokens = usage.get("completion_tokens") + if (pt := usage.get("prompt_tokens")) is not None: + output.prompt_len = pt + if first_chunk_received: + output.success = True + else: + output.success = False + output.error = ( + "Never received a valid chunk to calculate TTFT." + "This response will be marked as failed!" + ) + output.generated_text = generated_text + output.latency = most_recent_timestamp - st + 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 + + +def _get_chat_content( + request_func_input: RequestFuncInput, + mm_position: Literal["first", "last"] = "last", +) -> list[dict[str, Any]]: + text_contents = [{"type": "text", "text": request_func_input.prompt}] + + mm_contents = [] + if request_func_input.multi_modal_content: + mm_content = request_func_input.multi_modal_content + if isinstance(mm_content, list): + mm_contents.extend(request_func_input.multi_modal_content) + elif isinstance(mm_content, dict): + mm_contents.append(request_func_input.multi_modal_content) + else: + raise TypeError( + "multi_modal_content must be a dict or list[dict] for openai-chat" + ) + + if mm_position == "first": + return mm_contents + text_contents + + return text_contents + mm_contents + + +async def async_request_openai_chat_completions( + request_func_input: RequestFuncInput, + session: aiohttp.ClientSession, + pbar: tqdm | None = None, + mm_position: Literal["first", "last"] = "last", +) -> RequestFuncOutput: + api_url = request_func_input.api_url + _validate_api_url(api_url, "OpenAI Chat Completions API", "chat/completions") + + content = _get_chat_content(request_func_input, mm_position=mm_position) + + payload = { + "model": ( + request_func_input.model_name + if request_func_input.model_name + else request_func_input.model + ), + "messages": [ + {"role": "user", "content": content}, + ], + "max_completion_tokens": request_func_input.output_len, + "stream": True, + "stream_options": { + "include_usage": True, + }, + } + _update_payload_common(payload, request_func_input) + + headers = _get_headers("application/json") + _update_headers_common(headers, request_func_input) + + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + generated_text = "" + ttft = 0.0 + st = time.perf_counter() + output.start_time = st + most_recent_timestamp = st + try: + async with session.post(url=api_url, json=payload, headers=headers) as response: + if response.status == 200: + handler = StreamedResponseHandler() + async for chunk_bytes in response.content.iter_any(): + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + messages = handler.add_chunk(chunk_bytes) + for message in messages: + if message.startswith(":"): + continue + + chunk = message.removeprefix("data: ") + + if chunk != "[DONE]": + timestamp = time.perf_counter() + data = json.loads(chunk) + + if choices := data.get("choices"): + content = choices[0]["delta"].get("content") + if ttft == 0.0: + ttft = timestamp - st + output.ttft = ttft + else: + output.itl.append(timestamp - most_recent_timestamp) + + generated_text += content or "" + elif usage := data.get("usage"): + output.output_tokens = usage.get("completion_tokens") + if (pt := usage.get("prompt_tokens")) is not None: + output.prompt_len = pt + + most_recent_timestamp = timestamp + + output.generated_text = generated_text + output.success = True + output.latency = most_recent_timestamp - st + 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 = { + "openai": async_request_openai_completions, + "tokenspeed": async_request_openai_completions, + "openai-chat": async_request_openai_chat_completions, +} + + +def get_model(pretrained_model_name_or_path: str) -> str: + if envs.TOKENSPEED_USE_MODELSCOPE.get(): + import huggingface_hub.constants + from modelscope import snapshot_download + + return 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 pretrained_model_name_or_path + + +def get_tokenizer( + pretrained_model_name_or_path: str, +) -> PreTrainedTokenizerBase: + 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 download_and_cache_file(url: str, filename: str | None = None) -> str: + if filename is None: + filename = os.path.join("/tmp", os.path.basename(urlparse(url).path)) + if os.path.exists(filename): + return filename + + print(f"Downloading from {url} to {filename}") + response = requests.get(url, stream=True) + response.raise_for_status() + total_size = int(response.headers.get("content-length", 0)) + 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=1024): + f.write(chunk) + bar.update(len(chunk)) + return filename + + +def is_valid_sequence( + prompt_len: int, + output_len: int, + max_model_len: int | None, + skip_min_tokens_check: bool, +) -> bool: + if not skip_min_tokens_check and (prompt_len < 4 or output_len < 4): + return False + if max_model_len is not None and prompt_len + output_len > max_model_len: + return False + return True + + +def _resolve_range_ratios( + range_ratio: RangeRatio, +) -> tuple[float, float]: + """Return ``(input_range_ratio, output_range_ratio)`` from *range_ratio*. + + *range_ratio* is either a single float (used for both input and output) + or a dict with ``"input"`` and ``"output"`` keys. + """ + if isinstance(range_ratio, dict): + try: + return float(range_ratio["input"]), float(range_ratio["output"]) + except KeyError as exc: + raise ValueError( + "When range_ratio is a dict it must contain 'input' and " + f"'output' keys, got: {sorted(range_ratio)}" + ) from exc + ratio = float(range_ratio) + return ratio, ratio + + +def get_sampling_params( + rng: np.random.Generator, + num_requests: int, + range_ratio: RangeRatio, + input_len: int, + output_len: int, + tokenizer: PreTrainedTokenizerBase, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """ + Sample per-request input/output token lengths and vocab offsets. + + Lengths are drawn uniformly from integer ranges around the configured + means, controlled by *range_ratio*. It may be a single ``float`` + (applied to both input and output) or a ``dict`` with ``"input"`` and + ``"output"`` keys for independent control. + + Tokenizer special tokens are subtracted from ``input_len`` before + computing the sampling interval. + + Returns: + (input_lens, output_lens, offsets) - three 1-D ``np.ndarray`` of + shape ``(num_requests,)``. + """ + input_range_ratio, output_range_ratio = _resolve_range_ratios(range_ratio) + + if not (0.0 <= input_range_ratio < 1.0): + raise ValueError("input_range_ratio must be in [0, 1).") + if not (0.0 <= output_range_ratio < 1.0): + raise ValueError("output_range_ratio must be in [0, 1).") + num_special_tokens = int(tokenizer.num_special_tokens_to_add()) + real_input_len = max(0, int(input_len) - num_special_tokens) + input_low = math.floor(real_input_len * (1 - input_range_ratio)) + input_high = math.ceil(real_input_len * (1 + input_range_ratio)) + output_low = math.floor(output_len * (1 - output_range_ratio)) + output_high = math.ceil(output_len * (1 + output_range_ratio)) + # Ensure the lower bound for output length is at least 1 to + # prevent sampling 0 tokens. + output_low = max(output_low, 1) + output_high = max(output_high, 1) + + if input_low > input_high: + raise ValueError( + f"Invalid input sampling interval: low={input_low} > high={input_high}" + ) + if output_low > output_high: + raise ValueError( + f"Invalid output sampling interval: low={output_low} > high={output_high}" + ) + + logger.info( + "Sampling input_len from [%s, %s] and output_len from [%s, %s]", + input_low, + input_high, + output_low, + output_high, + ) + + input_lens = rng.integers(input_low, input_high + 1, size=num_requests) + output_lens = rng.integers(output_low, output_high + 1, size=num_requests) + offsets = rng.integers(0, tokenizer.vocab_size, size=num_requests) + return input_lens, output_lens, offsets + + +def gen_prompt_decode_to_target_len( + tokenizer: PreTrainedTokenizerBase, + token_sequence: list[int], + target_token_len: int, + max_retry: int = 10, + add_special_tokens: bool = False, + rng: np.random.Generator | None = None, +) -> tuple[str, list[int], int]: + """ + Ensure decoded-then-encoded prompt length matches the target token length. + + This function decodes an initial token sequence to text and re-encodes it + , iteratively adjusting the token sequence length to match a target. + This is necessary because some tokenizers do not guarantee a 1:1 mapping + between consecutive tokens and the decoded-then-encoded sequence length. + For example, for GPT2Tokenizer: + [6880, 6881] -> ['Ġcalls', 'here'] -> + [1650, 939, 486] -> ['Ġcall', 'sh', 'ere'] + + Returns a tuple of the final prompt string, the adjusted token sequence, + and the token mismatch (final_len - target_token_len) if the retry budget + is exhausted. + """ + remain_num_try = max_retry + token_mismatch = 0 + while True: + prompt = tokenizer.decode(token_sequence) + token_sequence = tokenizer.encode(prompt, add_special_tokens=add_special_tokens) + if remain_num_try <= 0: + if len(token_sequence) != target_token_len: + token_mismatch = len(token_sequence) - target_token_len + break + + if len(token_sequence) == target_token_len: + break + elif len(token_sequence) < target_token_len: + if rng is not None: + extra_tokens = rng.integers( + 0, + tokenizer.vocab_size, + size=target_token_len - len(token_sequence), + ).tolist() + else: + extra_tokens = np.random.randint( + 0, + tokenizer.vocab_size, + size=target_token_len - len(token_sequence), + ).tolist() + token_sequence.extend(extra_tokens) + elif len(token_sequence) > target_token_len: + token_sequence = token_sequence[:target_token_len] + + remain_num_try -= 1 + + return prompt, token_sequence, token_mismatch + + +class BenchmarkDataset: + DEFAULT_SEED = 0 + + def __init__( + self, + dataset_path: str | None = None, + random_seed: int = DEFAULT_SEED, + disable_shuffle: bool = False, + **kwargs, + ) -> None: + """ + Initialize the BenchmarkDataset with an optional dataset path and random + seed. + """ + self.dataset_path = dataset_path + self.random_seed = random_seed if random_seed is not None else self.DEFAULT_SEED + self.disable_shuffle = disable_shuffle + self.data: Any | None = None + + def get_lora_request( + self, + index: int, + max_loras: int | None = None, + lora_path: str | None = None, + lora_assignment: str = "random", + ) -> None: + return None + + +# fmt: off +class RandomDataset(BenchmarkDataset): + """ + Synthetic text-only dataset for serving/throughput benchmarks. + + Strategy: + - Sample input/output token lengths per request from integer-uniform ranges + around configured means (controlled by range_ratio). + - Prepend a fixed random prefix of length prefix_len. + - Generate the remaining tokens as a reproducible sequence: + (offset + index + arange(input_len)) % vocab_size. + - Decode then re-encode/truncate to ensure prompt token counts match. + - Uses numpy.default_rng seeded with random_seed for reproducible sampling. + """ + + DEFAULT_PREFIX_LEN = 0 + DEFAULT_RANGE_RATIO = 0.0 + DEFAULT_INPUT_LEN = 1024 + DEFAULT_OUTPUT_LEN = 128 + + def __init__(self, **kwargs) -> None: + super().__init__(**kwargs) + # Use numpy's default_rng for deterministic sampling + # Do not use random.seed() or np.random.seed() elsewhere in this class. + # This ensures that the RNG is isolated from global RNG state. + self._rng = np.random.default_rng(self.random_seed) + + def sample( + self, + tokenizer: PreTrainedTokenizerBase, + num_requests: int, + request_id_prefix: str = "", + no_oversample: bool = False, + prefix_len: int = DEFAULT_PREFIX_LEN, + range_ratio: RangeRatio = DEFAULT_RANGE_RATIO, + input_len: int = DEFAULT_INPUT_LEN, + output_len: int = DEFAULT_OUTPUT_LEN, + batchsize: int = 1, + max_loras: int | None = None, + lora_path: str | None = None, + lora_assignment: str = "random", + **kwargs, + ) -> list[SampleRequest]: + resolved_input_rr, _ = _resolve_range_ratios(range_ratio) + + num_special = int(tokenizer.num_special_tokens_to_add()) + real_input_len = max(0, int(input_len) - num_special) + min_sampled_input = math.floor( + real_input_len * (1.0 - float(resolved_input_rr)) + ) + min_total_input = int(prefix_len) + min_sampled_input + if min_total_input < 1: + raise ValueError( + "--random-input-len is too small: with tokenizer special " + f"tokens {num_special} and " + f"input range ratio {resolved_input_rr}, " + "the minimum possible total input tokens (prefix + sampled) is " + f"{min_total_input}. Increase --random-input-len and/or " + "--random-prefix-len, or decrease the input range ratio " + "so that prefix_len + floor(max(0, random_input_len - " + "num_special)) * (1 - input_range_ratio) >= 1." + ) + + input_lens, output_lens, offsets = get_sampling_params( + self._rng, + num_requests, + range_ratio, + input_len, + output_len, + tokenizer, + ) + + vocab_size = tokenizer.vocab_size + prohibited_tokens = tokenizer.all_special_ids + all_tokens = np.arange(vocab_size) + allowed_tokens = np.array(list(set(all_tokens) - set(prohibited_tokens))) + + # Generate prefix once + prefix_token_ids = self.get_prefix(tokenizer, allowed_tokens, prefix_len) + + requests = [] + token_mismatch_total = 0 + for i in range(num_requests): + prompt, total_input_len, token_mismatch = self.generate_token_sequence( # noqa: E501 + tokenizer=tokenizer, + prefix_token_ids=prefix_token_ids, + prefix_len=prefix_len, + vocab_size=vocab_size, + input_len=int(input_lens[i]), + offset=int(offsets[i]), + index=i, + allowed_tokens=allowed_tokens, + ) + token_mismatch_total += token_mismatch + lora_req = self.get_lora_request( + index=i, + max_loras=max_loras, + lora_path=lora_path, + lora_assignment=lora_assignment, + ) + requests.append( + SampleRequest( + prompt=prompt, + prompt_len=total_input_len, + expected_output_len=int(output_lens[i]), + lora_request=lora_req, + request_id=request_id_prefix + str(i), + ) + ) + # only used for embeddings benchmark. + if batchsize > 1: + batch_requests = [] + # Create batched requests + for i in range(0, num_requests, batchsize): + batch = requests[i : i + batchsize] + batch_requests.append( + SampleRequest( + prompt=[req.prompt for req in batch], + prompt_len=sum(req.prompt_len for req in batch), + expected_output_len=0, + request_id=request_id_prefix + str(i // batchsize), + ) + ) + requests = batch_requests + + if token_mismatch_total != 0: + sign = "more" if token_mismatch_total > 0 else "fewer" + logger.warning( + "Across all generated prompts, there were %d %s tokens " + "than expected after decoding and re-encoding. This is " + "expected due to the imperfect nature of the sampling " + "procedure.", + abs(token_mismatch_total), + sign, + ) + + return requests + + def get_prefix( + self, + tokenizer: PreTrainedTokenizerBase, + allowed_tokens: np.ndarray, + prefix_len: int, + ) -> list[int]: + """ + Get the prefix for the dataset. + """ + if prefix_len <= 0: + return [] + + prefix_tokens = allowed_tokens[ + self._rng.integers(0, len(allowed_tokens), size=prefix_len) + ].tolist() + _, adjusted_tokens, token_mismatch = gen_prompt_decode_to_target_len( + tokenizer=tokenizer, + token_sequence=prefix_tokens, + target_token_len=prefix_len, + add_special_tokens=False, + rng=self._rng, + ) + if token_mismatch != 0: + sign = "more" if token_mismatch > 0 else "fewer" + logger.warning( + "Prefix tokenization produced %d %s tokens than expected " + "after decoding and re-encoding. This is expected due to " + "the imperfect nature of the sampling procedure", + abs(token_mismatch), + sign, + ) + return adjusted_tokens + + def generate_token_sequence( + self, + *, + tokenizer: PreTrainedTokenizerBase, + prefix_token_ids: list[int], + prefix_len: int, + vocab_size: int, + input_len: int, + offset: int, + index: int, + allowed_tokens: np.ndarray, + ) -> tuple[str, int, int]: + """ + Returns (prompt, total_input_len). + + NOTE: After decoding the prompt we have to encode and decode it again. + This is done because in some cases N consecutive tokens + give a string tokenized into != N number of tokens. + For example for GPT2Tokenizer: + [6880, 6881] -> ['Ġcalls', 'here'] -> + [1650, 939, 486] -> ['Ġcall', 'sh', 'ere'] + To avoid uncontrolled change of the prompt length, + the encoded sequence is truncated before being decoded again. + """ + # Build the inner sequence by sampling + # sequentially from the allowed tokens + inner_seq = allowed_tokens[ + (offset + index + np.arange(input_len)) % len(allowed_tokens) + ].tolist() + token_sequence = prefix_token_ids + inner_seq + + # Decode, then re-encode and truncate to preserve token count invariants + total_input_len = prefix_len + int(input_len) + prompt, adjusted_token_sequence, token_mismatch = ( + gen_prompt_decode_to_target_len( + tokenizer=tokenizer, + token_sequence=token_sequence, + target_token_len=total_input_len, + add_special_tokens=False, + rng=self._rng, + ) + ) + total_input_len = len(adjusted_token_sequence) + return prompt, total_input_len, token_mismatch +# fmt: on + + +def sample_sharegpt_requests( + dataset_path: str | None, + num_requests: int, + tokenizer: PreTrainedTokenizerBase, + fixed_output_len: int | None = None, + max_model_len: int | None = None, + apply_chat_template: bool = False, + skip_min_tokens_check: bool = False, +) -> list[SampleRequest]: + if fixed_output_len is not None and fixed_output_len < 4: + raise ValueError("output_len too small") + if not dataset_path: + dataset_path = download_and_cache_file(SHAREGPT_URL) + + with open(dataset_path, encoding="utf-8") as f: + dataset = json.load(f) + + conversations = [] + for data in dataset: + turns = data.get("conversations", data.get("conversation", [])) + if len(turns) >= 2: + conversations.append((turns[0]["value"], turns[1]["value"])) + random.shuffle(conversations) + + samples: list[SampleRequest] = [] + for prompt, completion in conversations: + if len(samples) == num_requests: + break + if apply_chat_template: + prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + add_generation_prompt=True, + tokenize=False, + ) + if tokenizer.bos_token: + prompt = prompt.replace(tokenizer.bos_token, "") + prompt_len = len(tokenizer.encode(prompt)) + output_len = fixed_output_len or len(tokenizer.encode(completion)) + if not is_valid_sequence( + prompt_len, output_len, max_model_len, skip_min_tokens_check + ): + continue + samples.append(SampleRequest(prompt, prompt_len, output_len)) + + print(f"#Input tokens: {sum(x.prompt_len for x in samples)}") + print(f"#Output tokens: {sum(x.expected_output_len for x in samples)}") + return samples + + +def sample_random_requests( + input_len: int, + output_len: int, + num_prompts: int, + range_ratio: float, + tokenizer: PreTrainedTokenizerBase, + dataset_path: str | None, + prefix_len: int = 0, + random_seed: int = 0, + request_id_prefix: str = "", +) -> list[SampleRequest]: + if dataset_path is not None: + raise ValueError("Cannot use 'random' dataset with --dataset-path.") + + samples = RandomDataset(random_seed=random_seed).sample( + tokenizer=tokenizer, + num_requests=num_prompts, + request_id_prefix=request_id_prefix, + prefix_len=prefix_len, + range_ratio=range_ratio, + input_len=input_len, + output_len=output_len, + ) + + print(f"#Input tokens: {sum(x.prompt_len for x in samples)}") + print(f"#Output tokens: {sum(x.expected_output_len for x in samples)}") + return samples + + +def get_samples( + args: argparse.Namespace, tokenizer: PreTrainedTokenizerBase +) -> list[SampleRequest]: + if args.dataset_name == "sharegpt": + return sample_sharegpt_requests( + dataset_path=args.dataset_path, + num_requests=args.num_prompts, + tokenizer=tokenizer, + fixed_output_len=args.sharegpt_output_len, + max_model_len=args.max_model_len, + apply_chat_template=args.apply_chat_template, + skip_min_tokens_check=args.skip_min_tokens_check, + ) + if args.dataset_name == "random": + return 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, + prefix_len=args.random_prefix_len, + random_seed=args.seed, + request_id_prefix=args.request_id_prefix, + ) + raise ValueError(f"Unknown dataset: {args.dataset_name}") + + +def get_current_request_rate( + ramp_up_strategy: Literal["linear", "exponential"] | None, + ramp_up_start_rps: int | None, + ramp_up_end_rps: int | None, + request_index: int, + total_requests: int, + request_rate: float, +) -> float: + if ( + ramp_up_strategy + and ramp_up_start_rps is not None + and ramp_up_end_rps is not None + ): + progress = request_index / max(total_requests - 1, 1) + if ramp_up_strategy == "linear": + return ramp_up_start_rps + (ramp_up_end_rps - ramp_up_start_rps) * progress + if ramp_up_strategy == "exponential": + ratio = ramp_up_end_rps / ramp_up_start_rps + return ramp_up_start_rps * (ratio**progress) + raise ValueError(f"Unknown ramp-up strategy: {ramp_up_strategy}") + return request_rate + + +async def get_request( + input_requests: list[SampleRequest], + request_rate: float, + burstiness: float = 1.0, + ramp_up_strategy: Literal["linear", "exponential"] | None = None, + ramp_up_start_rps: int | None = None, + ramp_up_end_rps: int | None = None, +) -> AsyncGenerator[tuple[SampleRequest, float], None]: + assert ( + burstiness > 0 + ), f"A positive burstiness factor is expected, got {burstiness}." + total_requests = len(input_requests) + assert total_requests > 0, "No requests provided." + + delay_ts = [] + request_rates = [] + for request_index, _request in enumerate(input_requests): + current_request_rate = get_current_request_rate( + ramp_up_strategy, + ramp_up_start_rps, + ramp_up_end_rps, + request_index, + total_requests, + request_rate, + ) + assert ( + current_request_rate > 0.0 + ), f"Non-positive request rate {current_request_rate}." + request_rates.append(current_request_rate) + if current_request_rate == float("inf"): + delay_ts.append(0.0) + elif burstiness == float("inf"): + delay_ts.append(1.0 / current_request_rate) + else: + theta = 1.0 / (current_request_rate * burstiness) + delay_ts.append(float(np.random.gamma(shape=burstiness, scale=theta))) + + for i in range(1, len(delay_ts)): + delay_ts[i] += delay_ts[i - 1] + if ramp_up_strategy is None and delay_ts[-1] != 0: + target_total_delay_s = total_requests / request_rate + normalize_factor = target_total_delay_s / delay_ts[-1] + delay_ts = [delay * normalize_factor for delay in delay_ts] + + start_ts = time.time() + for request_index, request in enumerate(input_requests): + if delay_ts[request_index] > 0: + sleep_interval_s = start_ts + delay_ts[request_index] - time.time() + if sleep_interval_s > 0: + await asyncio.sleep(sleep_interval_s) + yield request, request_rates[request_index] + + +async def get_first_model_from_server( + base_url: str, + headers: dict[str, str] | None = None, + ssl_context: ssl.SSLContext | bool | None = None, +) -> tuple[str, str]: + connector = aiohttp.TCPConnector(ssl=ssl_context) + async with aiohttp.ClientSession(connector=connector) as session: + models_url = f"{base_url}/v1/models" + try: + async with session.get(models_url, headers=headers) as response: + response.raise_for_status() + data = await response.json() + if data.get("data"): + model = data["data"][0] + return model["id"], model.get("root", model["id"]) + raise ValueError(f"No models found on the server at {base_url}.") + except (aiohttp.ClientError, json.JSONDecodeError) as e: + raise RuntimeError(f"Failed to fetch models from {models_url}: {e}") from e + + +async def wait_for_endpoint( + request_func, + test_input: RequestFuncInput, + session: aiohttp.ClientSession, + timeout_seconds: int = 600, + retry_interval: int = 5, +) -> RequestFuncOutput: + deadline = time.perf_counter() + timeout_seconds + output = RequestFuncOutput(success=False) + print(f"Waiting for endpoint to become up in {timeout_seconds} seconds") + with tqdm( + total=timeout_seconds, + bar_format="{desc} |{bar}| {elapsed} elapsed, {remaining} remaining", + unit="s", + ) as pbar: + while True: + remaining = deadline - time.perf_counter() + elapsed = timeout_seconds - remaining + pbar.update(min(elapsed - pbar.n, timeout_seconds - pbar.n)) + pbar.refresh() + if remaining <= 0: + break + try: + output = await request_func(test_input, session=session) + if output.success: + return output + err_last_line = str(output.error).rstrip().rsplit("\n", 1)[-1] + print(f"Endpoint is not ready. Error='{err_last_line}'") + except aiohttp.ClientConnectorError: + pass + await asyncio.sleep(min(retry_interval, max(remaining, 0))) + return output + + +def calculate_metrics( + input_requests: list[SampleRequest], + outputs: list[RequestFuncOutput], + dur_s: float, + tokenizer: PreTrainedTokenizerBase | None, + selected_percentiles: list[float], + goodput_config_dict: dict[str, float], +) -> tuple[BenchmarkMetrics, list[int]]: + actual_output_lens: list[int] = [] + total_input = 0 + completed = 0 + good_completed = 0 + itls: list[float] = [] + tpots: list[float] = [] + all_tpots: list[float] = [] + ttfts: list[float] = [] + e2els: list[float] = [] + + for output in outputs: + if output.success: + output_len = output.output_tokens + if not output_len: + output_len = ( + len( + tokenizer.encode( + output.generated_text, add_special_tokens=False + ) + ) + if tokenizer + else 1 + ) + actual_output_lens.append(output_len) + total_input += output.prompt_len + tpot = 0.0 + if output_len > 1: + tpot = (output.latency - output.ttft) / (output_len - 1) + tpots.append(tpot) + all_tpots.append(tpot) + itls.extend(output.itl) + ttfts.append(output.ttft) + e2els.append(output.latency) + completed += 1 + else: + actual_output_lens.append(0) + + if goodput_config_dict: + valid_metrics = [] + slo_values = [] + if "ttft" in goodput_config_dict: + valid_metrics.append(ttfts) + slo_values.append( + goodput_config_dict["ttft"] / MILLISECONDS_TO_SECONDS_CONVERSION + ) + if "tpot" in goodput_config_dict: + valid_metrics.append(all_tpots) + slo_values.append( + goodput_config_dict["tpot"] / MILLISECONDS_TO_SECONDS_CONVERSION + ) + if "e2el" in goodput_config_dict: + valid_metrics.append(e2els) + slo_values.append( + goodput_config_dict["e2el"] / MILLISECONDS_TO_SECONDS_CONVERSION + ) + for req_metric in zip(*valid_metrics): + if all(slo >= metric for slo, metric in zip(slo_values, req_metric)): + good_completed += 1 + + if completed == 0: + warnings.warn( + "All requests failed. This is likely due to a misconfiguration on the benchmark arguments.", + stacklevel=2, + ) + + successful_outputs = [output for output in outputs if output.success] + failed_outputs = [output for output in outputs if not output.success] + if failed_outputs: + print("Failed requests during benchmark run detected (capping to 10):") + for i, err in enumerate(failed_outputs[:10]): + print(f"Error {i}: {err.error}") + + max_output_tokens_per_s = 0.0 + max_concurrent_requests = 0 + if successful_outputs: + min_start_time = min(output.start_time for output in successful_outputs) + max_end_time = max( + output.start_time + output.latency for output in successful_outputs + ) + duration_seconds = int(np.ceil(max_end_time - min_start_time)) + 1 + tokens_per_second = np.zeros(duration_seconds) + concurrent_requests_per_second = np.zeros(duration_seconds) + for output in successful_outputs: + token_times = [output.start_time + output.ttft] + current_time = token_times[0] + for itl_value in output.itl: + current_time += itl_value + token_times.append(current_time) + for token_time in token_times: + second_bucket = int(token_time - min_start_time) + if 0 <= second_bucket < duration_seconds: + tokens_per_second[second_bucket] += 1 + request_start_second = int(output.start_time - min_start_time) + request_end_second = int( + (output.start_time + output.latency) - min_start_time + ) + for second in range(request_start_second, request_end_second + 1): + concurrent_requests_per_second[second] += 1 + max_output_tokens_per_s = ( + float(np.max(tokens_per_second)) if len(tokens_per_second) else 0.0 + ) + max_concurrent_requests = ( + int(np.max(concurrent_requests_per_second)) + if len(concurrent_requests_per_second) + else 0 + ) + + metrics = BenchmarkMetrics( + completed=completed, + failed=len(failed_outputs), + total_input=total_input, + total_output=sum(actual_output_lens), + request_throughput=completed / dur_s, + request_goodput=good_completed / dur_s, + output_throughput=sum(actual_output_lens) / dur_s, + total_token_throughput=(total_input + sum(actual_output_lens)) / dur_s, + mean_ttft_ms=np.mean(ttfts or 0) * 1000, + std_ttft_ms=np.std(ttfts or 0) * 1000, + median_ttft_ms=np.median(ttfts or 0) * 1000, + percentiles_ttft_ms=[ + (p, np.percentile(ttfts or 0, p) * 1000) for p in selected_percentiles + ], + mean_tpot_ms=np.mean(tpots or 0) * 1000, + std_tpot_ms=np.std(tpots or 0) * 1000, + median_tpot_ms=np.median(tpots or 0) * 1000, + percentiles_tpot_ms=[ + (p, np.percentile(tpots or 0, p) * 1000) for p in selected_percentiles + ], + mean_itl_ms=np.mean(itls or 0) * 1000, + std_itl_ms=np.std(itls or 0) * 1000, + median_itl_ms=np.median(itls or 0) * 1000, + percentiles_itl_ms=[ + (p, np.percentile(itls or 0, p) * 1000) for p in selected_percentiles + ], + mean_e2el_ms=np.mean(e2els or 0) * 1000, + std_e2el_ms=np.std(e2els or 0) * 1000, + median_e2el_ms=np.median(e2els or 0) * 1000, + percentiles_e2el_ms=[ + (p, np.percentile(e2els or 0, p) * 1000) for p in selected_percentiles + ], + max_output_tokens_per_s=max_output_tokens_per_s, + max_concurrent_requests=max_concurrent_requests, + ) + return metrics, actual_output_lens + + +async def benchmark( + task_type: TaskType, + backend: str, + api_url: str, + base_url: str, + model_id: str, + model_name: str | None, + tokenizer: PreTrainedTokenizerBase | None, + input_requests: list[SampleRequest], + logprobs: int | None, + request_rate: float, + burstiness: float, + disable_tqdm: bool, + num_warmups: int, + profile: bool, + profile_num_steps: int | None, + selected_percentile_metrics: list[str], + selected_percentiles: list[float], + ignore_eos: bool, + goodput_config_dict: dict[str, float], + max_concurrency: int | None, + extra_headers: dict[str, str] | None, + extra_body: dict[str, Any] | None, + ramp_up_strategy: Literal["linear", "exponential"] | None = None, + ramp_up_start_rps: int | None = None, + ramp_up_end_rps: int | None = None, + ready_check_timeout_sec: int = 600, + ssl_context: ssl.SSLContext | bool | None = None, +) -> dict[str, Any]: + try: + request_func = ASYNC_REQUEST_FUNCS[backend] + except KeyError: + raise ValueError(f"Unknown backend: {backend}") from None + + connector = aiohttp.TCPConnector(ssl=ssl_context) + session = aiohttp.ClientSession( + connector=connector, trust_env=True, timeout=AIOHTTP_TIMEOUT + ) + + test_request = input_requests[0] + test_input = RequestFuncInput( + model=model_id, + model_name=model_name, + prompt=test_request.prompt, + api_url=api_url, + prompt_len=test_request.prompt_len, + output_len=test_request.expected_output_len, + logprobs=logprobs, + ignore_eos=ignore_eos, + extra_headers=extra_headers, + extra_body=extra_body, + ) + + if ready_check_timeout_sec > 0: + print("Starting initial single prompt test run...") + test_output = await wait_for_endpoint( + request_func, test_input, session, timeout_seconds=ready_check_timeout_sec + ) + if not test_output.success: + raise ValueError( + "Initial test run failed - Please make sure benchmark arguments are correctly specified. " + f"Error: {test_output.error}" + ) + print("Initial test run completed.") + else: + print("Skipping endpoint ready check.") + + if num_warmups > 0: + print(f"Warming up with {num_warmups} requests...") + warmup_pbar = None if disable_tqdm else tqdm(total=num_warmups) + warmup_semaphore = ( + asyncio.Semaphore(max_concurrency) + if max_concurrency + else contextlib.nullcontext() + ) + + async def warmup_limited_request_func(): + async with warmup_semaphore: + return await request_func(test_input, session=session, pbar=warmup_pbar) + + await asyncio.gather( + *( + asyncio.create_task(warmup_limited_request_func()) + for _ in range(num_warmups) + ) + ) + if warmup_pbar: + warmup_pbar.close() + print("Warmup run completed.") + + if profile: + if profile_num_steps is None: + print("Starting profiler...") + else: + print(f"Starting profiler for {profile_num_steps} steps...") + extra_body = dict(extra_body or {}) + if profile_num_steps is not None: + extra_body["num_steps"] = profile_num_steps + profile_input = RequestFuncInput( + model=model_id, + model_name=model_name, + prompt=test_request.prompt, + api_url=base_url + "/start_profile", + prompt_len=test_request.prompt_len, + output_len=test_request.expected_output_len, + logprobs=logprobs, + ignore_eos=ignore_eos, + extra_headers=extra_headers, + extra_body=extra_body, + ) + profile_output = await request_func(profile_input, session=session) + if profile_output.success: + print("Profiler started") + + distribution = "Poisson process" if burstiness == 1.0 else "Gamma distribution" + if ramp_up_strategy: + print(f"Traffic ramp-up strategy: {ramp_up_strategy}.") + print(f"Will increase RPS from {ramp_up_start_rps} to {ramp_up_end_rps} RPS.") + else: + print(f"Traffic request rate: {request_rate}") + print(f"Burstiness factor: {burstiness} ({distribution})") + print(f"Maximum request concurrency: {max_concurrency}") + + pbar = None if disable_tqdm else tqdm(total=len(input_requests)) + semaphore = ( + asyncio.Semaphore(max_concurrency) + if max_concurrency + else contextlib.nullcontext() + ) + + async def limited_request_func(request_func_input, session, pbar): + async with semaphore: + coro = request_func(request_func_input, session=session, pbar=pbar) + return await await_with_per_request_timeout( + coro, + prompt_len=request_func_input.prompt_len, + pbar=pbar, + ) + + print("Starting main benchmark run...") + benchmark_start_time = time.perf_counter() + tasks: list[asyncio.Task] = [] + rps_change_events = [] + last_int_rps = -1 + if ramp_up_strategy is not None and ramp_up_start_rps is not None: + last_int_rps = ramp_up_start_rps + rps_change_events.append( + {"rps": last_int_rps, "timestamp": datetime.now().isoformat()} + ) + + async for request, current_request_rate in get_request( + input_requests, + request_rate, + burstiness, + ramp_up_strategy, + ramp_up_start_rps, + ramp_up_end_rps, + ): + if ramp_up_strategy is not None: + current_int_rps = int(current_request_rate) + if current_int_rps > last_int_rps: + timestamp = datetime.now().isoformat() + for rps_val in range(last_int_rps + 1, current_int_rps + 1): + rps_change_events.append({"rps": rps_val, "timestamp": timestamp}) + last_int_rps = current_int_rps + request_func_input = RequestFuncInput( + model=model_id, + model_name=model_name, + prompt=request.prompt, + api_url=api_url, + prompt_len=request.prompt_len, + output_len=request.expected_output_len, + logprobs=logprobs, + ignore_eos=ignore_eos, + extra_headers=extra_headers, + extra_body=extra_body, + request_id=request.request_id, + ) + tasks.append( + asyncio.create_task(limited_request_func(request_func_input, session, pbar)) + ) + + outputs = await asyncio.gather(*tasks) + if pbar: + pbar.close() + benchmark_duration = time.perf_counter() - benchmark_start_time + + metrics, actual_output_lens = calculate_metrics( + input_requests, + outputs, + benchmark_duration, + tokenizer, + selected_percentiles, + goodput_config_dict, + ) + + _print_section_header(" Serving Benchmark Result ", "=") + _print_metric_row("Successful requests:", metrics.completed) + _print_metric_row("Failed requests:", metrics.failed) + if max_concurrency is not None: + _print_metric_row("Maximum request concurrency:", max_concurrency) + if request_rate != float("inf"): + _print_metric_row("Request rate configured (RPS):", request_rate, precision=2) + _print_metric_row("Benchmark duration (s):", benchmark_duration, precision=2) + _print_metric_row("Total input tokens:", metrics.total_input) + _print_metric_row("Total generated tokens:", metrics.total_output) + _print_metric_row( + "Request throughput (req/s):", metrics.request_throughput, precision=2 + ) + if goodput_config_dict: + _print_metric_row( + "Request goodput (req/s):", metrics.request_goodput, precision=2 + ) + _print_metric_row( + "Output token throughput (tok/s):", metrics.output_throughput, precision=2 + ) + _print_metric_row( + "Peak output token throughput (tok/s):", + metrics.max_output_tokens_per_s, + precision=2, + ) + _print_metric_row( + "Peak concurrent requests:", metrics.max_concurrent_requests, precision=2 + ) + _print_metric_row( + "Total token throughput (tok/s):", + metrics.total_token_throughput, + precision=2, + ) + + result: dict[str, Any] = { + "duration": benchmark_duration, + "completed": metrics.completed, + "failed": metrics.failed, + "total_input_tokens": metrics.total_input, + "total_output_tokens": metrics.total_output, + "request_throughput": metrics.request_throughput, + "request_goodput": metrics.request_goodput if goodput_config_dict else None, + "output_throughput": metrics.output_throughput, + "total_token_throughput": metrics.total_token_throughput, + "input_lens": [output.prompt_len for output in outputs], + "output_lens": actual_output_lens, + "ttfts": [output.ttft for output in outputs], + "itls": [output.itl for output in outputs], + "start_times": [output.start_time for output in outputs], + "generated_texts": [output.generated_text for output in outputs], + "errors": [output.error for output in outputs], + "max_output_tokens_per_s": metrics.max_output_tokens_per_s, + "max_concurrent_requests": metrics.max_concurrent_requests, + } + if rps_change_events: + result["rps_change_events"] = rps_change_events + + def process_one_metric( + metric_attribute_name: str, metric_name: str, metric_header: str + ) -> None: + if metric_attribute_name not in selected_percentile_metrics: + return + _print_section_header(metric_header, "-") + _print_metric_row( + f"Mean {metric_name} (ms):", + getattr(metrics, f"mean_{metric_attribute_name}_ms"), + precision=2, + ) + _print_metric_row( + f"Median {metric_name} (ms):", + getattr(metrics, f"median_{metric_attribute_name}_ms"), + precision=2, + ) + result[f"mean_{metric_attribute_name}_ms"] = getattr( + metrics, f"mean_{metric_attribute_name}_ms" + ) + result[f"median_{metric_attribute_name}_ms"] = getattr( + metrics, f"median_{metric_attribute_name}_ms" + ) + result[f"std_{metric_attribute_name}_ms"] = getattr( + metrics, f"std_{metric_attribute_name}_ms" + ) + for p, value in getattr(metrics, f"percentiles_{metric_attribute_name}_ms"): + p_word = str(int(p)) if int(p) == p else str(p) + _print_metric_row(f"P{p_word} {metric_name} (ms):", value, precision=2) + result[f"p{p_word}_{metric_attribute_name}_ms"] = value + + process_one_metric("ttft", "TTFT", "Time to First Token") + process_one_metric("tpot", "TPOT", "Time per Output Token (excl. 1st token)") + process_one_metric("itl", "ITL", "Inter-token Latency") + process_one_metric("e2el", "E2EL", "End-to-end Latency") + + print("=" * 50) + + if profile and profile_num_steps is None: + print("Stopping profiler...") + profile_input = RequestFuncInput( + model=model_id, + model_name=model_name, + prompt=test_request.prompt, + api_url=base_url + "/stop_profile", + prompt_len=test_request.prompt_len, + output_len=test_request.expected_output_len, + logprobs=logprobs, + ignore_eos=ignore_eos, + ) + profile_output = await request_func(profile_input, session=session) + if profile_output.success: + print("Profiler stopped") + + await session.close() + return result + + +def parse_goodput(slo_pairs: list[str] | None) -> dict[str, float]: + goodput_config_dict: dict[str, float] = {} + if not slo_pairs: + return goodput_config_dict + try: + for slo_pair in slo_pairs: + slo_name, slo_val = slo_pair.split(":") + goodput_config_dict[slo_name] = float(slo_val) + except ValueError as err: + raise argparse.ArgumentTypeError( + 'Specify service level objectives for goodput as "KEY:VALUE" pairs.' + ) from err + for slo_name, slo_val in goodput_config_dict.items(): + if slo_name not in {"ttft", "tpot", "e2el"}: + raise ValueError(f"Invalid goodput metric {slo_name!r}.") + if slo_val < 0: + raise ValueError(f"Goodput SLO {slo_name!r} must be non-negative.") + return goodput_config_dict + + +def compute_result_filename( + args: argparse.Namespace, model_id: str, label: str | None, current_dt: str +) -> str | None: + if not (args.save_result or args.append_result or args.output_file): + return None + if args.output_file: + return args.output_file + base_model_id = model_id.split("/")[-1] + max_concurrency_str = ( + f"-concurrency{args.max_concurrency}" + if args.max_concurrency is not None + else "" + ) + result_label = label or args.backend + if args.ramp_up_strategy is not None: + file_name = f"{result_label}-ramp-up-{args.ramp_up_strategy}-{args.ramp_up_start_rps}qps-{args.ramp_up_end_rps}qps{max_concurrency_str}-{base_model_id}-{current_dt}.json" + else: + file_name = f"{result_label}-{args.request_rate}qps{max_concurrency_str}-{base_model_id}-{current_dt}.json" + if args.result_dir: + os.makedirs(args.result_dir, exist_ok=True) + file_name = os.path.join(args.result_dir, file_name) + return file_name + + +def add_dataset_parser(parser: argparse.ArgumentParser) -> None: + parser.add_argument( + "--dataset-name", + type=str, + default="random", + choices=["sharegpt", "random"], + help="Name of the dataset to benchmark on.", + ) + parser.add_argument( + "--dataset-path", type=str, default=None, help="Path to the dataset." + ) + parser.add_argument("--num-prompts", type=int, default=DEFAULT_NUM_PROMPTS) + parser.add_argument("--input-len", type=int, default=None) + parser.add_argument("--output-len", type=int, default=None) + parser.add_argument("--max-model-len", type=int, default=None) + parser.add_argument("--skip-min-tokens-check", action="store_true") + parser.add_argument("--sharegpt-output-len", type=int, default=None) + parser.add_argument("--random-input-len", type=int, default=1024) + parser.add_argument("--random-output-len", type=int, default=128) + parser.add_argument("--random-range-ratio", type=float, default=0.0) + parser.add_argument("--random-prefix-len", type=int, default=0) + parser.add_argument("--request-id-prefix", type=str, default="bench-") + + +def add_serving_cli_args(parser: argparse.ArgumentParser) -> None: + add_dataset_parser(parser) + parser.add_argument("--label", type=str, default=None) + parser.add_argument( + "--backend", + type=str, + default="openai", + choices=list(ASYNC_REQUEST_FUNCS.keys()), + help="The backend type to use for the benchmark.", + ) + parser.add_argument("--base-url", type=str, default=None) + parser.add_argument("--host", type=str, default="127.0.0.1") + parser.add_argument("--port", type=int, default=8000) + parser.add_argument("--endpoint", type=str, default="/v1/completions") + parser.add_argument("--header", metavar="KEY=VALUE", nargs="*") + parser.add_argument("--model", type=str, default=None) + parser.add_argument("--served-model-name", type=str, default=None) + parser.add_argument("--tokenizer", type=str, default=None) + parser.add_argument("--skip-tokenizer-init", action="store_true") + parser.add_argument("--trust-remote-code", action="store_true", default=True) + parser.add_argument("--request-rate", type=float, default=float("inf")) + parser.add_argument("--burstiness", type=float, default=1.0) + parser.add_argument("--max-concurrency", type=int, default=None) + parser.add_argument("--num-warmups", type=int, default=0) + parser.add_argument("--ready-check-timeout-sec", type=int, default=600) + parser.add_argument("--disable-tqdm", action="store_true") + parser.add_argument("--profile", action="store_true") + parser.add_argument("--profile-num-steps", type=int, default=None) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--ignore-eos", action="store_true") + parser.add_argument("--disable-ignore-eos", action="store_true") + parser.add_argument("--apply-chat-template", action="store_true") + parser.add_argument("--logprobs", type=int, default=None) + parser.add_argument("--extra-body", type=json.loads, default={}) + parser.add_argument("--extra-request-body", type=json.loads, default=None) + parser.add_argument("--goodput", nargs="*", default=None) + parser.add_argument("--percentile-metrics", type=str, default=None) + parser.add_argument("--metric-percentiles", type=str, default="99") + parser.add_argument( + "--ramp-up-strategy", choices=["linear", "exponential"], default=None + ) + parser.add_argument("--ramp-up-start-rps", type=int, default=None) + parser.add_argument("--ramp-up-end-rps", type=int, default=None) + parser.add_argument("--insecure", action="store_true") + parser.add_argument("--save-result", action="store_true") + parser.add_argument("--append-result", action="store_true") + parser.add_argument("--save-detailed", action="store_true") + parser.add_argument("--result-dir", type=str, default=None) + parser.add_argument("--output-file", type=str, default=None) + parser.set_defaults(dispatch_function=BenchmarkServingSubcommand.cmd) + + +async def main_async(args: argparse.Namespace) -> dict[str, Any]: + print(args) + set_ulimit() + random.seed(args.seed) + np.random.seed(args.seed) + + if args.disable_ignore_eos: + args.ignore_eos = False + if args.extra_request_body is not None: + args.extra_body = args.extra_request_body + if args.profile_num_steps is not None: + if args.profile_num_steps <= 0: + raise ValueError("--profile-num-steps must be positive.") + if not args.profile: + raise ValueError("--profile-num-steps requires --profile.") + if args.input_len is not None: + args.random_input_len = args.input_len + if args.output_len is not None: + args.random_output_len = args.output_len + args.sharegpt_output_len = args.output_len + + if args.ramp_up_strategy is not None: + if args.request_rate != float("inf"): + raise ValueError("When using ramp-up, do not specify --request-rate.") + if args.ramp_up_start_rps is None or args.ramp_up_end_rps is None: + raise ValueError( + "Ramp-up requires --ramp-up-start-rps and --ramp-up-end-rps." + ) + if args.ramp_up_start_rps > args.ramp_up_end_rps: + raise ValueError("Ramp-up start RPS must be less than end RPS.") + if args.ramp_up_strategy == "exponential" and args.ramp_up_start_rps == 0: + raise ValueError("For exponential ramp-up, start RPS cannot be 0.") + + if args.base_url is not None: + api_url = f"{args.base_url}{args.endpoint}" + base_url = args.base_url + else: + host_port = join_host_port(args.host, args.port) + api_url = f"http://{host_port}{args.endpoint}" + base_url = f"http://{host_port}" + + headers = None + if args.header: + headers = {} + for item in args.header: + if "=" not in item: + raise ValueError("Invalid header format. Please use KEY=VALUE format.") + key, value = item.split("=", 1) + headers[key.strip()] = value.strip() + + ssl_context: ssl.SSLContext | bool | None = ( + False if args.insecure else True if base_url.startswith("https://") else None + ) + + if args.model is None: + print("Model not specified, fetching first model from server...") + model_name, model_id = await get_first_model_from_server( + base_url, headers, ssl_context + ) + print(f"First model name: {model_name}, first model id: {model_id}") + else: + model_name = args.served_model_name + model_id = args.model + + tokenizer = None + tokenizer_id = None + if not args.skip_tokenizer_init: + tokenizer_id = args.tokenizer or model_id + tokenizer = get_tokenizer(tokenizer_id) + + if args.dataset_name == "random" and args.backend in OPENAI_COMPATIBLE_BACKENDS: + args.ignore_eos = True + + input_requests = get_samples(args, tokenizer) + goodput_config_dict = parse_goodput(args.goodput) + extra_body = args.extra_body or {} + percentile_metrics = args.percentile_metrics or "ttft,tpot,itl" + + if "temperature" not in extra_body: + print( + "WARNING: tokenspeed bench serve no longer sets temperature==0 in requests by default. " + "The server decides its own default. Include --extra-body '{\"temperature\": 0}' for greedy decoding." + ) + + benchmark_result = await benchmark( + task_type=TaskType.GENERATION, + backend=args.backend, + api_url=api_url, + base_url=base_url, + model_id=model_id, + model_name=model_name, + tokenizer=tokenizer, + input_requests=input_requests, + logprobs=args.logprobs, + request_rate=args.request_rate, + burstiness=args.burstiness, + disable_tqdm=args.disable_tqdm, + num_warmups=args.num_warmups, + profile=args.profile, + profile_num_steps=args.profile_num_steps, + selected_percentile_metrics=percentile_metrics.split(","), + selected_percentiles=[float(p) for p in args.metric_percentiles.split(",")], + ignore_eos=args.ignore_eos, + goodput_config_dict=goodput_config_dict, + max_concurrency=args.max_concurrency, + extra_headers=headers, + extra_body=extra_body, + ramp_up_strategy=args.ramp_up_strategy, + ramp_up_start_rps=args.ramp_up_start_rps, + ramp_up_end_rps=args.ramp_up_end_rps, + ready_check_timeout_sec=args.ready_check_timeout_sec, + ssl_context=ssl_context, + ) + + current_dt = datetime.now().strftime("%Y%m%d-%H%M%S") + result_json = { + "date": current_dt, + "backend": args.backend, + "label": args.label, + "model_id": model_id, + "tokenizer_id": tokenizer_id, + "num_prompts": args.num_prompts, + "request_rate": ( + args.request_rate if args.request_rate < float("inf") else "inf" + ), + "burstiness": args.burstiness, + "max_concurrency": args.max_concurrency, + **benchmark_result, + } + + if not args.save_detailed: + for field_name in [ + "input_lens", + "output_lens", + "start_times", + "ttfts", + "itls", + "generated_texts", + "errors", + ]: + result_json.pop(field_name, None) + + file_name = compute_result_filename(args, model_id, args.label, current_dt) + if file_name: + with open( + file_name, mode="a+" if args.append_result else "w", encoding="utf-8" + ) as outfile: + if args.append_result and outfile.tell() != 0: + outfile.write("\n") + json.dump(result_json, outfile) + + return result_json + + +def run_benchmark(args: argparse.Namespace) -> dict[str, Any]: + return asyncio.run(main_async(args)) + + +class BenchmarkSubcommandBase: + help: str + name: str + + @classmethod + def add_cli_args(cls, parser: argparse.ArgumentParser) -> None: + raise NotImplementedError + + @staticmethod + def cmd(args: argparse.Namespace) -> None: + raise NotImplementedError + + +class BenchmarkServingSubcommand(BenchmarkSubcommandBase): + name = "serve" + help = "Benchmark online serving throughput." + + @classmethod + def add_cli_args(cls, parser: argparse.ArgumentParser) -> None: + add_serving_cli_args(parser) + + @staticmethod + def cmd(args: argparse.Namespace) -> None: + run_benchmark(args) + + +class BenchmarkSubcommand: + name = "bench" + help = "TokenSpeed bench subcommand." + + @staticmethod + def cmd(args: argparse.Namespace) -> None: + args.dispatch_function(args) + + def subparser_init( + self, subparsers: argparse._SubParsersAction + ) -> argparse.ArgumentParser: + bench_parser = subparsers.add_parser( + self.name, + help=self.help, + description=self.help, + usage=f"tokenspeed {self.name} [options]", + ) + bench_subparsers = bench_parser.add_subparsers(required=True, dest="bench_type") + for cmd_cls in BenchmarkSubcommandBase.__subclasses__(): + cmd_subparser = bench_subparsers.add_parser( + cmd_cls.name, + help=cmd_cls.help, + description=cmd_cls.help, + usage=f"tokenspeed {self.name} {cmd_cls.name} [options]", + ) + cmd_subparser.set_defaults(dispatch_function=cmd_cls.cmd) + cmd_cls.add_cli_args(cmd_subparser) + return bench_parser + + +def is_legacy_serving_args(argv: list[str]) -> bool: + return bool(argv) and argv[0].startswith("-") and argv[0] not in ("-h", "--help") + + +def main(argv: list[str] | None = None) -> None: + argv = list(sys.argv[1:] if argv is None else argv) + if is_legacy_serving_args(argv): + parser = argparse.ArgumentParser(description=BenchmarkServingSubcommand.help) + BenchmarkServingSubcommand.add_cli_args(parser) + args = parser.parse_args(argv) + BenchmarkServingSubcommand.cmd(args) + return + + parser = argparse.ArgumentParser( + prog="tokenspeed", description="TokenSpeed benchmark commands." + ) + subparsers = parser.add_subparsers(required=True, dest="command") + BenchmarkSubcommand().subparser_init(subparsers) + args = parser.parse_args(["bench", *argv]) + BenchmarkSubcommand.cmd(args) + + +if __name__ == "__main__": + main() diff --git a/python/tokenspeed/cli/__init__.py b/python/tokenspeed/cli/__init__.py new file mode 100644 index 0000000..7182ad2 --- /dev/null +++ b/python/tokenspeed/cli/__init__.py @@ -0,0 +1,25 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""TokenSpeed CLI package.""" + +from tokenspeed.cli.__main__ import main + +__all__ = ["main"] diff --git a/python/tokenspeed/cli/__main__.py b/python/tokenspeed/cli/__main__.py new file mode 100644 index 0000000..955e763 --- /dev/null +++ b/python/tokenspeed/cli/__main__.py @@ -0,0 +1,108 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""TokenSpeed CLI entry point.""" + +import argparse +import sys + + +def _serve(args: argparse.Namespace, raw_argv: list[str]) -> None: + from tokenspeed.cli.serve_smg import run_smg_from_args + + run_smg_from_args(args, raw_argv) + + +def _bench(args: argparse.Namespace) -> None: + from tokenspeed.bench import main as bench_main + + bench_main(args.bench_args) + + +def _env(args: argparse.Namespace) -> None: + from tokenspeed.env import main as env_main + + env_main() + + +def _version(args: argparse.Namespace) -> None: + from tokenspeed.version import __version__ + + print(f"TokenSpeed v{__version__}") + + +def main() -> None: + parser = argparse.ArgumentParser( + prog="tokenspeed", + description="TokenSpeed is a speed-of-light LLM inference engine.", + ) + + subparsers = parser.add_subparsers(dest="command") + + # Unknown flags fall through to the smg orchestrator's own splitter; we + # don't register the engine's ServerArgs on this parser. + serve_parser = subparsers.add_parser( + "serve", + help="Launch the TokenSpeed inference server.", + ) + serve_parser.set_defaults(func=_serve) + + bench_parser = subparsers.add_parser( + "bench", + add_help=False, + help="Run TokenSpeed benchmark commands.", + ) + bench_parser.set_defaults(func=_bench, bench_args=[]) + + env_parser = subparsers.add_parser( + "env", + help="Check environment configurations and dependency versions.", + ) + env_parser.set_defaults(func=_env) + + version_parser = subparsers.add_parser( + "version", + help="Print the TokenSpeed version.", + ) + version_parser.set_defaults(func=_version) + + args, extra_args = parser.parse_known_args() + + if args.command is None: + parser.print_help() + sys.exit(1) + + if args.func is _bench: + args.bench_args = extra_args + args.func(args) + return + + if args.func is _serve: + raw = list(sys.argv[2:]) + args.func(args, raw) + return + + if extra_args: + parser.error(f"unrecognized arguments: {' '.join(extra_args)}") + args.func(args) + + +if __name__ == "__main__": + main() diff --git a/python/tokenspeed/cli/_argsplit.py b/python/tokenspeed/cli/_argsplit.py new file mode 100644 index 0000000..52b1926 --- /dev/null +++ b/python/tokenspeed/cli/_argsplit.py @@ -0,0 +1,243 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Argv splitter for ``ts serve``. + +A leading positional argument is treated as the model (vllm-style +``ts serve [flags...]``) and rewritten to ``--model `` +before routing. + +Routing precedence is top-down. The first matching rule wins: + +1. Orchestrator-only flags (consumed, never forwarded) +2. ``--model`` / ``--reasoning-parser`` — fanned out to both. + ``--reasoning-parser`` goes to the gateway (post-gen parsing) and + the engine (defers JSON grammars past the reasoning channel). +3. ``--host`` / ``--port`` — gateway only (user-facing) +4. ``--chat-template`` / ``--tool-call-parser`` — gateway only + **(override)**: ``prepare_server_args`` accepts these too, but in smg + mode the gateway owns OpenAI-compat HTTP and parsing. +5. ``--tp`` / ``--tensor-parallel-size`` — engine only (alias normalized) +6. Anything else ``prepare_server_args`` accepts — engine only +7. Anything else — gateway (fall-through to ``smg launch`` clap) +""" + +from __future__ import annotations + +import argparse +import functools +from dataclasses import dataclass, field +from typing import Iterable + +_ORCH_FLAGS = { + "--engine-startup-timeout", + "--gateway-startup-timeout", + "--drain-timeout", + "--control-port", +} + +_FANOUT_FLAGS = {"--model", "--reasoning-parser"} + +_ALIASES = { + "--model-path": "--model", + "--tp": "--tensor-parallel-size", +} + +_GATEWAY_USER_FACING = {"--host", "--port"} + +_GATEWAY_OVERRIDE = { + "--chat-template", + "--tool-call-parser", +} + +_ENGINE_EXPLICIT = {"--tensor-parallel-size"} + +_MODEL_FLAG_TOKENS = ("--model", "--model-path") + +_ENGINE_MULTI_VALUE_FLAGS = { + "--cudagraph-capture-sizes", +} + + +def _has_model_flag(tokens: Iterable[str]) -> bool: + for token in tokens: + if token in _MODEL_FLAG_TOKENS: + return True + for flag in _MODEL_FLAG_TOKENS: + if token.startswith(flag + "="): + return True + return False + + +@dataclass +class OrchestratorOpts: + engine_startup_timeout: int = 1800 + gateway_startup_timeout: int = 60 + drain_timeout: int = 30 + control_port: int | None = None + + +@dataclass +class SplitResult: + engine: list[str] = field(default_factory=list) + gateway: list[str] = field(default_factory=list) + opts: OrchestratorOpts = field(default_factory=OrchestratorOpts) + + +def _normalize(argv: Iterable[str]) -> list[tuple[str, list[str] | str | None]]: + """Convert raw argv into a list of (name, value) pairs. + + Handles both ``--flag value`` and ``--flag=value`` forms. Aliases are + resolved to their canonical names. + """ + items: list[tuple[str, list[str] | str | None]] = [] + tokens = list(argv) + i = 0 + while i < len(tokens): + raw = tokens[i] + if not raw.startswith("--"): + raise ValueError(f"unexpected positional arg: {raw!r}") + if "=" in raw: + name, _, value = raw.partition("=") + name = _ALIASES.get(name, name) + i += 1 + else: + name = _ALIASES.get(raw, raw) + nxt = tokens[i + 1] if i + 1 < len(tokens) else None + if nxt is None or nxt.startswith("--"): + value = None + i += 1 + elif name in _ENGINE_MULTI_VALUE_FLAGS: + values = [] + i += 1 + while i < len(tokens) and not tokens[i].startswith("--"): + values.append(tokens[i]) + i += 1 + value = values + else: + value = nxt + i += 2 + items.append((name, value)) + return items + + +def _append_arg(args: list[str], name: str, value: list[str] | str | None) -> None: + if value is None: + args.append(name) + elif isinstance(value, list): + args.extend([name, *value]) + else: + args.extend([name, value]) + + +@functools.lru_cache(maxsize=1) +def _engine_recognized_flags() -> set[str]: + """Snapshot the set of long-form flags accepted by ``prepare_server_args``.""" + # Lazy import: ServerArgs pulls the full runtime stack (~200ms). + from tokenspeed.runtime.utils.server_args import ServerArgs + + parser = argparse.ArgumentParser(add_help=False) + ServerArgs.add_cli_args(parser) + flags: set[str] = set() + for action in parser._actions: + for opt in action.option_strings: + if opt.startswith("--"): + flags.add(opt) + flags.discard("--help") + flags.discard("-h") + return flags + + +def split_argv(argv: list[str]) -> SplitResult: + """Split ts-serve argv into engine_args, gateway_args, orchestrator_opts. + + A leading positional argument is rewritten to ``--model `` so + ``ts serve [flags...]`` and ``ts serve --model [flags...]`` + both work. + + Raises: + ValueError: if a flag that requires a value is provided without one + (e.g. ``--model`` with no path), if a timeout flag is + non-positive, if the model is given both positionally and via + ``--model``/``--model-path``, or if a positional arg appears + after the leading model. + """ + + argv = list(argv) + if argv and not argv[0].startswith("--"): + model = argv[0] + rest = argv[1:] + if _has_model_flag(rest): + raise ValueError( + "model specified both as positional argument and via " + "--model/--model-path" + ) + argv = ["--model", model, *rest] + + items = _normalize(argv) + result = SplitResult() + engine_flags = _engine_recognized_flags() + + for name, value in items: + if name in _ORCH_FLAGS: + if value is None or value == "": + raise ValueError(f"{name} requires a positive integer (seconds)") + try: + seconds = int(value) + except ValueError as e: + raise ValueError(f"{name}={value!r} is not a valid integer") from e + if seconds <= 0: + raise ValueError(f"{name} must be positive, got {seconds}") + attr = name[2:].replace("-", "_") + setattr(result.opts, attr, seconds) + continue + + if name in _FANOUT_FLAGS: + if value is None: + raise ValueError(f"{name} requires a value") + _append_arg(result.engine, name, value) + _append_arg(result.gateway, name, value) + continue + + if name in _GATEWAY_USER_FACING: + if value is None: + raise ValueError(f"{name} requires a value") + _append_arg(result.gateway, name, value) + continue + + if name in _GATEWAY_OVERRIDE: + if value is None: + raise ValueError(f"{name} requires a value") + _append_arg(result.gateway, name, value) + continue + + if name in _ENGINE_EXPLICIT: + if value is None: + raise ValueError(f"{name} requires a value") + _append_arg(result.engine, name, value) + continue + + if name in engine_flags: + _append_arg(result.engine, name, value) + continue + + _append_arg(result.gateway, name, value) + + return result diff --git a/python/tokenspeed/cli/_logo.py b/python/tokenspeed/cli/_logo.py new file mode 100644 index 0000000..91091c8 --- /dev/null +++ b/python/tokenspeed/cli/_logo.py @@ -0,0 +1,100 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Startup banner for ``ts serve``: ``Token`` in bright black, ``Speed`` in blue.""" + +from __future__ import annotations + +import os +import sys +from typing import IO + +from tokenspeed.version import __version__ + +_TOKEN_LINES = ( + "████████ ██████ ██ ██ ███████ ███ ██ ", + " ██ ██ ██ ██ ██ ██ ████ ██ ", + " ██ ██ ██ █████ █████ ██ ██ ██ ", + " ██ ██ ██ ██ ██ ██ ██ ██ ██ ", + " ██ ██████ ██ ██ ███████ ██ ████ ", +) + +_SPEED_LINES = ( + "███████ ██████ ███████ ███████ ██████ ", + "██ ██ ██ ██ ██ ██ ██ ", + "███████ ██████ █████ █████ ██ ██ ", + " ██ ██ ██ ██ ██ ██ ", + "███████ ██ ███████ ███████ ██████ ", +) + +_TOKEN_STYLE = "\033[90m" # bright black (grey). +_SPEED_STYLE = "\033[94m" # bright blue. +_RESET = "\033[0m" + +_TAGLINE = "Tokens at the speed of light" +_DISABLE_ENV = "TOKENSPEED_DISABLE_LOGO" +# Visible width of one banner row: len(token) + " " + len(speed). +_BANNER_WIDTH = len(_TOKEN_LINES[0]) + 1 + len(_SPEED_LINES[0]) + + +def _is_disabled() -> bool: + value = os.environ.get(_DISABLE_ENV, "").strip().lower() + return value not in ("", "0", "false", "no") + + +def _stream_supports_color(stream: IO[str]) -> bool: + isatty = getattr(stream, "isatty", None) + if not callable(isatty): + return False + try: + return bool(isatty()) + except ValueError: + return False + + +def render_logo(version: str = __version__, *, color: bool = True) -> str: + """Return the multi-line banner. Trailing newline included.""" + if color: + token_open, speed_open, close = _TOKEN_STYLE, _SPEED_STYLE, _RESET + else: + token_open = speed_open = close = "" + + rows = [ + f"{token_open}{token}{close} {speed_open}{speed}{close}" + for token, speed in zip(_TOKEN_LINES, _SPEED_LINES) + ] + footer = f"{_TAGLINE} · v{version}" + pad = max(0, (_BANNER_WIDTH - len(footer)) // 2) + rows.append(" " * pad + footer) + return "\n".join(rows) + "\n" + + +def print_logo(version: str = __version__, *, stream: IO[str] | None = None) -> None: + """Write the banner to ``stream`` (default stderr). + + No-op when ``TOKENSPEED_DISABLE_LOGO`` is truthy. ANSI colors are emitted + only when the target stream is a TTY. + """ + if _is_disabled(): + return + if stream is None: + stream = sys.stderr + stream.write(render_logo(version, color=_stream_supports_color(stream))) + stream.flush() diff --git a/python/tokenspeed/cli/_logprefix.py b/python/tokenspeed/cli/_logprefix.py new file mode 100644 index 0000000..fc5c6f5 --- /dev/null +++ b/python/tokenspeed/cli/_logprefix.py @@ -0,0 +1,52 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Async per-line stream prefixer for ``ts serve``.""" + +from __future__ import annotations + +import asyncio +from typing import Protocol + +ENGINE_TAG = "ts" +GATEWAY_TAG = "smg" + + +class _Sink(Protocol): + def write(self, data: str) -> int: ... + def flush(self) -> None: ... + + +async def tag_stream(reader: asyncio.StreamReader, tag: str, sink: _Sink) -> None: + """Read lines from ``reader`` until EOF, write ``[tag] `` to ``sink``. + + A trailing partial line (without a newline) is still emitted with a + synthesized newline so the last line of a crash message is not lost. + """ + prefix = f"[{tag}] " + while True: + line = await reader.readline() + if not line: + return + text = line.decode("utf-8", errors="replace") + if text.endswith("\n"): + sink.write(prefix + text) + else: + sink.write(prefix + text + "\n") diff --git a/python/tokenspeed/cli/_proc.py b/python/tokenspeed/cli/_proc.py new file mode 100644 index 0000000..96a1450 --- /dev/null +++ b/python/tokenspeed/cli/_proc.py @@ -0,0 +1,163 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Subprocess supervision helpers for ``ts serve``.""" + +from __future__ import annotations + +import asyncio +import contextlib +import logging +import os +import sys +import time + +logger = logging.getLogger(__name__) + +_GATEWAY_MODULE = "smg" +_ENGINE_MODULE_DEFAULT = "smg_grpc_servicer.tokenspeed" +_PIPE_LINE_LIMIT = 64 * 1024 * 1024 + + +async def spawn_engine( + args: list[str], + *, + host: str, + port: int, +) -> asyncio.subprocess.Process: + """Spawn the engine subprocess with PIPE stdio.""" + module = os.environ.get("TS_SERVE_ENGINE_MODULE", _ENGINE_MODULE_DEFAULT) + cmd = [ + sys.executable, + "-m", + module, + "--host", + host, + "--port", + str(port), + *args, + ] + logger.info("spawn engine: %s", " ".join(cmd)) + return await asyncio.create_subprocess_exec( + *cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + limit=_PIPE_LINE_LIMIT, + ) + + +async def spawn_gateway( + args: list[str], + *, + engine_host: str, + engine_port: int, +) -> asyncio.subprocess.Process: + """Spawn ``python -m smg launch`` with PIPE stdio.""" + cmd = [ + sys.executable, + "-m", + _GATEWAY_MODULE, + "launch", + "--worker-urls", + f"grpc://{engine_host}:{engine_port}", + "--disable-retries", + "--disable-circuit-breaker", + *args, + ] + logger.info("spawn gateway: %s", " ".join(cmd)) + return await asyncio.create_subprocess_exec( + *cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + limit=_PIPE_LINE_LIMIT, + ) + + +async def wait_grpc_serving( + target: str, + *, + timeout: float, + poll_interval: float = 1.0, +) -> None: + """Poll ``grpc.health.v1.Health.Check`` on ``target`` until SERVING.""" + import grpc + from grpc_health.v1 import health_pb2, health_pb2_grpc + + deadline = time.monotonic() + timeout + last_err: Exception | None = None + async with grpc.aio.insecure_channel(target) as channel: + stub = health_pb2_grpc.HealthStub(channel) + while time.monotonic() < deadline: + try: + resp = await stub.Check( + health_pb2.HealthCheckRequest(service=""), + timeout=2.0, + ) + if resp.status == health_pb2.HealthCheckResponse.SERVING: + return + except Exception as exc: # noqa: BLE001 + last_err = exc + await asyncio.sleep(poll_interval) + detail = f" (last error: {last_err!r})" if last_err is not None else "" + raise TimeoutError( + f"engine never reached SERVING on {target} within {timeout:.0f}s{detail}" + ) + + +async def wait_http_ready( + url: str, + *, + timeout: float, + poll_interval: float = 1.0, +) -> None: + """Poll HTTP ``GET `` until 200, or raise ``TimeoutError``.""" + import aiohttp + + deadline = time.monotonic() + timeout + last_err: Exception | None = None + async with aiohttp.ClientSession() as session: + while time.monotonic() < deadline: + try: + async with session.get( + url, timeout=aiohttp.ClientTimeout(total=2.0) + ) as resp: + if resp.status == 200: + return + except Exception as exc: # noqa: BLE001 + last_err = exc + await asyncio.sleep(poll_interval) + detail = f" (last error: {last_err!r})" if last_err is not None else "" + raise TimeoutError( + f"gateway never reached 200 on {url} within {timeout:.0f}s{detail}" + ) + + +async def terminate_then_kill(proc, *, drain_timeout: float) -> None: + """SIGTERM, wait up to ``drain_timeout``, then SIGKILL if still alive.""" + if proc.returncode is not None: + return + proc.terminate() + try: + await asyncio.wait_for(proc.wait(), timeout=drain_timeout) + except asyncio.TimeoutError: + with contextlib.suppress(ProcessLookupError): + proc.kill() + with contextlib.suppress(Exception): + await proc.wait() diff --git a/python/tokenspeed/cli/serve_smg.py b/python/tokenspeed/cli/serve_smg.py new file mode 100644 index 0000000..2fbe0ac --- /dev/null +++ b/python/tokenspeed/cli/serve_smg.py @@ -0,0 +1,613 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""``ts serve`` orchestrator: spawn smg gateway + gRPC engine, tag logs, probe +readiness, and tear down gateway-first on shutdown.""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import logging +import os +import signal +import sys +from pathlib import Path + +from tokenspeed.cli._argsplit import OrchestratorOpts, split_argv +from tokenspeed.cli._logo import print_logo +from tokenspeed.cli._logprefix import ENGINE_TAG, GATEWAY_TAG, tag_stream +from tokenspeed.cli._proc import ( + spawn_engine, + spawn_gateway, + terminate_then_kill, + wait_grpc_serving, + wait_http_ready, +) +from tokenspeed.runtime.utils.network import get_free_port +from tokenspeed.runtime.utils.process import kill_process_tree + +logger = logging.getLogger(__name__) + +DEFAULT_GATEWAY_HOST = "0.0.0.0" +DEFAULT_GATEWAY_PORT = 8000 +DEFAULT_REASONING_PARSER = "passthrough" +DEEPSEEK_V4_REASONING_PARSER = "deepseek_v31" +DEEPSEEK_V4_TOOL_CALL_PARSER = "deepseek_v4" +GLM_REASONING_PARSER = "glm45" +GLM_TOOL_CALL_PARSER = "glm47_moe" +DEFAULT_SMG_LOG_LEVEL = "warn" +DEFAULT_SMG_PROMETHEUS_PORT = 8413 +# smg routing policy for ``ts serve``. Distinct from DEFAULT_REASONING_PARSER, +# which happens to share the "passthrough" string but configures an unrelated +# flag (--reasoning-parser). +DEFAULT_SMG_POLICY = "passthrough" +# smg reliability knobs we always want disabled when launched under +# ts serve. These are tokenspeed-internal defaults: not surfaced via +# the ts CLI, not routed through split_argv. +_DEFAULT_SMG_DISABLE_FLAGS = ( + "--disable-circuit-breaker", + "--disable-retries", +) + + +def _check_serve_extra_installed() -> None: + import importlib.util + + missing: list[str] = [] + if importlib.util.find_spec("smg") is None: + missing.append("tokenspeed-smg") + if importlib.util.find_spec("smg_grpc_servicer.tokenspeed.server") is None: + missing.append("tokenspeed-smg-grpc-servicer") + if missing: + sys.stderr.write( + "ts serve requires the bundled gateway packages, normally installed\n" + "as part of `tokenspeed`. Reinstall tokenspeed to restore them:\n\n" + " pip install --force-reinstall --no-deps tokenspeed\n\n" + "or install them explicitly:\n\n" + " pip install \\\n" + " tokenspeed-smg \\\n" + " tokenspeed-smg-grpc-servicer \\\n" + " tokenspeed-smg-grpc-proto\n\n" + f"Missing: {', '.join(missing)}\n" + ) + sys.exit(1) + + +def _user_host_port_from_gateway_args(gateway_args: list[str]) -> tuple[str, int]: + """Pull --host / --port out of the gateway-bound argv. + + Defaults match TokenSpeed's public serving endpoint. The argv MUST + be in canonical ``[--flag, value, ...]`` form as produced by + ``split_argv``; equals-form (``--port=8000``) is not handled here. + """ + host = DEFAULT_GATEWAY_HOST + port = DEFAULT_GATEWAY_PORT + it = iter(gateway_args) + for token in it: + if token == "--host": + host = next(it) + elif token == "--port": + port = int(next(it)) + return host, port + + +def _gateway_args_with_default_port(gateway_args: list[str]) -> list[str]: + if "--port" in gateway_args: + return gateway_args + return [*gateway_args, "--port", str(DEFAULT_GATEWAY_PORT)] + + +def _gateway_args_with_default_reasoning_parser(gateway_args: list[str]) -> list[str]: + if "--reasoning-parser" in gateway_args: + return gateway_args + return [*gateway_args, "--reasoning-parser", DEFAULT_REASONING_PARSER] + + +def _gateway_args_with_smg_disable_defaults(gateway_args: list[str]) -> list[str]: + """Append the smg reliability-disable switches if they are not already there.""" + result = list(gateway_args) + for flag in _DEFAULT_SMG_DISABLE_FLAGS: + if flag not in result: + result.append(flag) + return result + + +def _gateway_args_with_default_policy(gateway_args: list[str]) -> list[str]: + """Front smg's single backend with the ``passthrough`` routing policy. + + ``ts serve`` always orchestrates exactly one engine endpoint, so smg's binary + default (``cache_aware``) is pure overhead here: it runs the load-aware worker + monitor and subscribes to the engine's KV events (``SubscribeKvEvents``). + Against engines that predate that RPC the subscription surfaced as + ``NotImplementedError: Method not implemented`` (smg#1794). The ``passthrough`` + policy (smg#1797) forwards every request to the single healthy worker with no + load balancing, load monitoring, or KV-event subscription. + + Default-when-unset: an explicit operator ``--policy`` is preserved. + + NOTE: ``--policy`` is whitelisted by smg's clap ``value_parser`` — a gateway + that predates smg#1797 rejects ``passthrough`` and fails to start. This + injection therefore requires a bundled ``tokenspeed-smg`` that ships smg#1797; + the pin in ``python/pyproject.toml`` must be bumped to such a release in + lockstep with this default. + """ + if "--policy" in gateway_args: + return gateway_args + return [*gateway_args, "--policy", DEFAULT_SMG_POLICY] + + +_TOKENIZER_CACHE_FLAGS = ( + "--tokenizer-cache-enable-l0", + "--tokenizer-cache-enable-l1", +) + + +def _gateway_args_with_default_tokenizer_cache(gateway_args: list[str]) -> list[str]: + """Default smg tokenizer caches (L0 + L1) ON for gateway-fronted launches. + + For agentic / chat-completions traffic with a shared system prompt + history, + L1 prefix-caching at special-token boundaries cuts TTFT by ~30% (verified + end-to-end on mm25). smg's own clap defaults leave both layers OFF. + + Opt-out: operators can pass ``--no-tokenizer-cache-enable-l0`` and/or + ``--no-tokenizer-cache-enable-l1`` to ``ts serve``. The ``--no-`` form is + intercepted here (smg's clap doesn't accept it natively) and prevents the + positive injection for that layer. + """ + result = list(gateway_args) + for flag in _TOKENIZER_CACHE_FLAGS: + no_flag = "--no-" + flag[2:] + if no_flag in result: + # Operator opted out: strip the --no- marker (smg rejects it) + # and skip the positive injection for this layer. + while no_flag in result: + result.remove(no_flag) + continue + if flag not in result: + result.append(flag) + return result + + +def _gateway_args_with_default_log_level(gateway_args: list[str]) -> list[str]: + if "--log-level" in gateway_args: + return gateway_args + return [*gateway_args, "--log-level", DEFAULT_SMG_LOG_LEVEL] + + +def _gateway_args_with_default_prometheus_port(gateway_args: list[str]) -> list[str]: + """Pin the smg Prometheus exporter to ``DEFAULT_SMG_PROMETHEUS_PORT``. + + smg's own default (``29000``) collides easily when multiple ``ts serve`` + instances share a host or when a previous run hasn't released the + port yet — the gateway then exits early and the tokenizer + registration job never runs, surfacing later as + ``tokenizer_not_found`` on the first request. Pinning a tokenspeed- + specific default keeps the port stable for our deployments while + still allowing an explicit override. + """ + if "--prometheus-port" in gateway_args: + return gateway_args + return [*gateway_args, "--prometheus-port", str(DEFAULT_SMG_PROMETHEUS_PORT)] + + +def _user_model_id(gateway_args: list[str]) -> str | None: + """Return the value of ``--model`` from a split gateway argv, or ``None``.""" + try: + idx = gateway_args.index("--model") + except ValueError: + return None + if idx + 1 >= len(gateway_args): + return None + return gateway_args[idx + 1] + + +def _is_deepseek_v4_model(model_id: str | None) -> bool: + if not model_id: + return False + + normalized = model_id.lower().replace("_", "-") + compact = normalized.replace("-", "") + if "deepseek-v4" in normalized or "deepseekv4" in compact: + return True + + config_path = Path(model_id) / "config.json" + if not config_path.is_file(): + return False + try: + with config_path.open() as f: + config = json.load(f) + except (OSError, json.JSONDecodeError): + return False + if not isinstance(config, dict): + return False + architectures = config.get("architectures") or [] + return ( + config.get("model_type") == "deepseek_v4" + or "DeepseekV4ForCausalLM" in architectures + ) + + +def _is_glm_dsa_model(model_id: str | None) -> bool: + if not model_id: + return False + + normalized = model_id.lower().replace("_", "-") + compact = normalized.replace("-", "") + if "glm-5" in normalized or "glm5" in compact: + return True + + config_path = Path(model_id) / "config.json" + if not config_path.is_file(): + return False + try: + with config_path.open() as f: + config = json.load(f) + except (OSError, json.JSONDecodeError): + return False + if not isinstance(config, dict): + return False + architectures = config.get("architectures") or [] + return config.get("model_type") == "glm_moe_dsa" or any( + arch in {"GlmMoeDsaForCausalLM", "GlmMoeDsaForCausalLMNextN"} + for arch in architectures + ) + + +def _args_with_default_model_parsers( + engine_args: list[str], gateway_args: list[str] +) -> tuple[list[str], list[str]]: + """Apply model-family parser defaults before smg gateway defaults. + + Reasoning parser defaults must be visible to both processes: smg extracts + reasoning_content after generation, while the engine uses the same parser + name to defer json_schema grammars past the reasoning channel. + """ + model_id = _user_model_id(gateway_args) or _user_model_id(engine_args) + engine_result = list(engine_args) + gateway_result = list(gateway_args) + + if _is_deepseek_v4_model(model_id): + if ( + "--reasoning-parser" not in engine_result + and "--reasoning-parser" not in gateway_result + ): + engine_result.extend(["--reasoning-parser", DEEPSEEK_V4_REASONING_PARSER]) + gateway_result.extend(["--reasoning-parser", DEEPSEEK_V4_REASONING_PARSER]) + if "--tool-call-parser" not in gateway_result: + gateway_result.extend(["--tool-call-parser", DEEPSEEK_V4_TOOL_CALL_PARSER]) + + elif _is_glm_dsa_model(model_id): + if ( + "--reasoning-parser" not in engine_result + and "--reasoning-parser" not in gateway_result + ): + engine_result.extend(["--reasoning-parser", GLM_REASONING_PARSER]) + gateway_result.extend(["--reasoning-parser", GLM_REASONING_PARSER]) + if "--tool-call-parser" not in gateway_result: + gateway_result.extend(["--tool-call-parser", GLM_TOOL_CALL_PARSER]) + + return engine_result, gateway_result + + +def _prewarm_hf_tokenizer(model_id: str) -> None: + """Download tokenizer artifacts to the HF cache before the gateway boots. + + smg fires its ``AddTokenizer`` job asynchronously after the engine + reports SERVING. On fast runners (e.g. b300) the first eval request + can race that fetch and fail with ``tokenizer_not_found``. Pulling + tokenizer files into the HF cache up front keeps the registration + fast regardless of engine startup speed. + """ + if not model_id or os.path.isdir(model_id): + return + try: + from huggingface_hub import snapshot_download + except ImportError: + return + try: + snapshot_download( + repo_id=model_id, + allow_patterns=[ + "tokenizer*", + "special_tokens_map*", + "vocab*", + "merges*", + "*.json", + ], + ) + except Exception as exc: # noqa: BLE001 + logger.warning("HF tokenizer prewarm failed for %s: %s", model_id, exc) + + +def _gateway_args_with_defaults(gateway_args: list[str]) -> list[str]: + gateway_args = _gateway_args_with_default_port(gateway_args) + gateway_args = _gateway_args_with_default_reasoning_parser(gateway_args) + gateway_args = _gateway_args_with_smg_disable_defaults(gateway_args) + gateway_args = _gateway_args_with_default_policy(gateway_args) + gateway_args = _gateway_args_with_default_tokenizer_cache(gateway_args) + gateway_args = _gateway_args_with_default_log_level(gateway_args) + return _gateway_args_with_default_prometheus_port(gateway_args) + + +async def _start_control_server( + *, + gateway_url: str, + engine_grpc_addr: str, + host: str, + port: int, + timeout: float = 30.0, +) -> bool: + """Start the control HTTP server in a daemon thread and wait for it to bind. + + Runs uvicorn alongside smg without blocking the orchestrator event loop. + Returns True once the server is accepting connections, or False if it + failed to bind (e.g. the port is already in use) or did not come up within + ``timeout`` seconds. Non-fatal: the smg gateway runs independently. + """ + import threading + + from tokenspeed.runtime.entrypoints.http_server import build_server + + server = build_server( + gateway_url=gateway_url, + engine_grpc_addr=engine_grpc_addr, + host=host, + port=port, + ) + thread = threading.Thread(target=server.run, daemon=True, name="ts-http-server") + thread.start() + + # uvicorn sets `started = True` only after the socket is bound and serving. + loop = asyncio.get_running_loop() + start = loop.time() + deadline = start + timeout + while not server.started: + if not thread.is_alive(): + return False # uvicorn raised during startup (e.g. AddrInUse) + if loop.time() >= deadline: + return False + await asyncio.sleep(0.05) + logger.info("control server bound in %.2fs", loop.time() - start) + return True + + +async def _stream_to(proc, tag: str) -> None: + await asyncio.gather( + tag_stream(proc.stdout, tag, sys.stdout), + tag_stream(proc.stderr, tag, sys.stderr), + ) + + +async def _drain_log(task: asyncio.Task, timeout: float = 2.0) -> None: + try: + await asyncio.wait_for(task, timeout=timeout) + except (asyncio.TimeoutError, asyncio.CancelledError): + task.cancel() + + +class _ShutdownDuringStartup(Exception): + pass + + +class _ChildExitedDuringStartup(Exception): + pass + + +async def _probe_or_stop( + probe_coro, stop_event: asyncio.Event, *, proc=None, label: str = "" +): + """Race a readiness probe against the stop event and (optionally) the + subprocess's own exit. + + - probe success → return result + - stop event → raise ``_ShutdownDuringStartup`` + - proc exits → raise ``_ChildExitedDuringStartup`` with returncode + label + """ + probe_task = asyncio.create_task(probe_coro) + stop_task = asyncio.create_task(stop_event.wait()) + tasks = [probe_task, stop_task] + proc_task = None + if proc is not None: + proc_task = asyncio.create_task(proc.wait()) + tasks.append(proc_task) + done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED) + for t in pending: + t.cancel() + if proc_task is not None and proc_task in done: + rc = proc_task.result() + raise _ChildExitedDuringStartup( + f"{label} subprocess exited with rc={rc} during startup; " + f"see [{label}] log lines above for the cause" + ) + if stop_task in done: + raise _ShutdownDuringStartup() + return probe_task.result() + + +async def run_smg( + *, + engine_args: list[str], + gateway_args: list[str], + opts: OrchestratorOpts, + user_host: str, + user_port: int, + _stop_event: asyncio.Event | None = None, +) -> int: + """Lifecycle loop. Returns the orchestrator's exit code.""" + engine = None + gateway = None + engine_log: asyncio.Task | None = None + gateway_log: asyncio.Task | None = None + + # Install signal handlers before spawning any subprocess so a Ctrl-C + # during the readiness probe doesn't skip terminate_then_kill. + stop = _stop_event if _stop_event is not None else asyncio.Event() + loop = asyncio.get_running_loop() + for sig in (signal.SIGTERM, signal.SIGINT): + try: + loop.add_signal_handler(sig, stop.set) + except NotImplementedError: + pass # Windows: signal handlers via asyncio aren't supported. Out of scope. + + try: + engine_port = get_free_port() + + engine = await spawn_engine(engine_args, host="127.0.0.1", port=engine_port) + engine_log = asyncio.create_task(_stream_to(engine, ENGINE_TAG)) + + await _probe_or_stop( + wait_grpc_serving( + f"127.0.0.1:{engine_port}", + timeout=float(opts.engine_startup_timeout), + ), + stop, + proc=engine, + label=ENGINE_TAG, + ) + + gateway = await spawn_gateway( + gateway_args, engine_host="127.0.0.1", engine_port=engine_port + ) + gateway_log = asyncio.create_task(_stream_to(gateway, GATEWAY_TAG)) + + await _probe_or_stop( + wait_http_ready( + f"http://{user_host}:{user_port}/readiness", + timeout=float(opts.gateway_startup_timeout), + ), + stop, + proc=gateway, + label=GATEWAY_TAG, + ) + + sys.stdout.write(f"ts serve ready on http://{user_host}:{user_port}\n") + sys.stdout.flush() + + control_port = ( + opts.control_port if opts.control_port is not None else user_port + 1 + ) + control_ok = await _start_control_server( + gateway_url=f"http://{user_host}:{user_port}", + engine_grpc_addr=f"127.0.0.1:{engine_port}", + host=user_host, + port=control_port, + ) + if control_ok: + sys.stdout.write( + f"ts control server ready on http://{user_host}:{control_port}\n" + ) + else: + sys.stderr.write( + f"WARNING: ts control server failed to bind on " + f"http://{user_host}:{control_port} (port in use?); " + f"serving continues without it\n" + ) + sys.stdout.flush() + + engine_wait = asyncio.create_task(engine.wait()) + gateway_wait = asyncio.create_task(gateway.wait()) + stop_wait = asyncio.create_task(stop.wait()) + + done, pending = await asyncio.wait( + [engine_wait, gateway_wait, stop_wait], + return_when=asyncio.FIRST_COMPLETED, + ) + + for task in pending: + task.cancel() + + rc_engine = engine.returncode if engine.returncode is not None else 0 + rc_gateway = gateway.returncode if gateway.returncode is not None else 0 + if rc_engine != 0: + return rc_engine + if rc_gateway != 0: + return rc_gateway + return 0 + + except _ChildExitedDuringStartup as exc: + logger.error("startup failed: %s", exc) + return 1 + except _ShutdownDuringStartup: + logger.info("shutdown signal received during startup; exiting cleanly") + return 0 + except TimeoutError as exc: + logger.error("startup failed: %s", exc) + return 1 + except KeyboardInterrupt: + logger.info("interrupted; exiting cleanly") + return 0 + finally: + # Shutdown order: gateway first, then engine. + if gateway is not None: + await terminate_then_kill(gateway, drain_timeout=opts.drain_timeout) + if engine is not None: + await terminate_then_kill(engine, drain_timeout=opts.drain_timeout) + + drain_tasks = [ + _drain_log(t) for t in (engine_log, gateway_log) if t is not None + ] + if drain_tasks: + await asyncio.gather(*drain_tasks, return_exceptions=True) + + # Final reap: walk only the children of our two known subprocesses — + # never os.getpid(), which under pytest would walk the test runner's + # children and SIGKILL unrelated test fixtures. + for proc in (engine, gateway): + if proc is not None: + try: + kill_process_tree(proc.pid, include_parent=False) + except Exception: # noqa: BLE001 + pass + + +def run_smg_from_args(args: argparse.Namespace, raw_argv: list[str]) -> None: + """Entry point called from cli/__main__.py for ``ts serve``.""" + try: + import setproctitle + + setproctitle.setproctitle("ts-serve") + except ImportError: + pass + + print_logo() + + _check_serve_extra_installed() + split = split_argv(raw_argv) + engine_args, gateway_args = _args_with_default_model_parsers( + split.engine, split.gateway + ) + gateway_args = _gateway_args_with_defaults(gateway_args) + user_host, user_port = _user_host_port_from_gateway_args(gateway_args) + + model_id = _user_model_id(gateway_args) + if model_id is not None: + _prewarm_hf_tokenizer(model_id) + rc = asyncio.run( + run_smg( + engine_args=engine_args, + gateway_args=gateway_args, + opts=split.opts, + user_host=user_host, + user_port=user_port, + ) + ) + sys.exit(rc) diff --git a/python/tokenspeed/env.py b/python/tokenspeed/env.py new file mode 100755 index 0000000..e19c59d --- /dev/null +++ b/python/tokenspeed/env.py @@ -0,0 +1,351 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Check environment configurations and dependency versions.""" + +import importlib.metadata +import os +import resource +import subprocess +import sys +from collections import OrderedDict, defaultdict + +import torch + + +def is_cuda_build() -> bool: + return torch.version.cuda is not None + + +def is_rocm_build() -> bool: + return getattr(torch.version, "hip", None) is not None + + +# List of packages to check versions +PACKAGE_LIST = [ + "tokenspeed", + "aiohttp", + "apache-tvm-ffi", + "compressed-tensors", + "dill", + "einops", + "fastapi", + "flashinfer-cubin", + "flashinfer-python", + "hf_transfer", + "huggingface_hub", + "instanttensor", + "modelscope", + "msgspec", + "ninja", + "numpy", + "nvidia-cutlass-dsl", + "nvidia-cutlass-dsl-libs-cu13", + "nvidia-ml-py", + "nvtx", + "openai", + "openai-harmony", + "orjson", + "packaging", + "partial-json-parser", + "peft", + "pillow", + "prometheus-client", + "psutil", + "pybase64", + "pybind11", + "pydantic", + "py-spy", + "PyYAML", + "pytest-asyncio", + "python-multipart", + "pyzmq", + "requests", + "setproctitle", + "tiktoken", + "tokenspeed-deepep", + "tokenspeed-deepgemm", + "tokenspeed-fa3", + "tokenspeed-fa4", + "tokenspeed-fast-hadamard-transform", + "tokenspeed-flashmla", + "tokenspeed-iris", + "tokenspeed-kernel", + "tokenspeed-kernel-amd", + "tokenspeed-mla", + "tokenspeed-mooncake", + "tokenspeed-proton", + "tokenspeed-smg", + "tokenspeed-smg-grpc-proto", + "tokenspeed-smg-grpc-servicer", + "tokenspeed-triton", + "tokenspeed-triton-kernels", + "tokenspeed-trtllm-kernel", + "torch", + "torch_memory_saver", + "torchvision", + "tqdm", + "transformers", + "uv", + "uvicorn", + "uvloop", + "viztracer", + "xgrammar", +] + + +def get_package_versions(packages: list[str]) -> dict[str, str]: + """Get versions of specified packages.""" + versions = {} + for package in packages: + package_name = package.split("==")[0].split(">=")[0].split("<=")[0] + try: + versions[package_name] = importlib.metadata.version(package_name) + except importlib.metadata.PackageNotFoundError: + versions[package_name] = "Package Not Found" + return versions + + +def get_cuda_info() -> dict[str, object]: + """Get CUDA-related information if available.""" + if is_cuda_build(): + 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_rocm_build(): + 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 + return {} + + +def _get_gpu_info() -> dict[str, str]: + """Get information about available GPUs.""" + devices = defaultdict(list) + capabilities = defaultdict(list) + for device_index in range(torch.cuda.device_count()): + devices[torch.cuda.get_device_name(device_index)].append(str(device_index)) + capability = torch.cuda.get_device_capability(device_index) + capabilities[f"{capability[0]}.{capability[1]}"].append(str(device_index)) + + 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 = next(iter(capabilities.items())) + 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() -> dict[str, str | None]: + """Get CUDA version information.""" + if is_cuda_build(): + 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 + if is_rocm_build(): + from torch.utils.cpp_extension import 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 + return {"CUDA_HOME": ""} + + +def _get_nvcc_info() -> dict[str, str]: + """Get NVCC version information.""" + if is_cuda_build(): + from torch.utils.cpp_extension import CUDA_HOME + + if not CUDA_HOME: + return {"NVCC": "Not Available"} + + try: + nvcc = os.path.join(CUDA_HOME, "bin/nvcc") + nvcc_output = subprocess.check_output([nvcc, "-V"], text=True).strip() + return { + "NVCC": nvcc_output[ + nvcc_output.rfind("Cuda compilation tools") : nvcc_output.rfind( + "Build" + ) + ].strip() + } + except (OSError, subprocess.SubprocessError): + return {"NVCC": "Not Available"} + elif is_rocm_build(): + from torch.utils.cpp_extension import ROCM_HOME + + if not ROCM_HOME: + return {"HIPCC": "Not Available"} + + try: + hipcc = os.path.join(ROCM_HOME, "bin/hipcc") + hipcc_output = subprocess.check_output( + [hipcc, "--version"], text=True + ).strip() + return { + "HIPCC": hipcc_output[ + hipcc_output.rfind("HIP version") : hipcc_output.rfind("AMD clang") + ].strip() + } + except (OSError, subprocess.SubprocessError): + return {"HIPCC": "Not Available"} + else: + return {"NVCC": "Not Available"} + + +def _get_cuda_driver_version() -> dict[str, str]: + """Get CUDA driver version.""" + if is_cuda_build(): + try: + output = subprocess.check_output( + [ + "nvidia-smi", + "--query-gpu=driver_version", + "--format=csv,noheader,nounits", + ], + text=True, + ) + versions = set(output.strip().splitlines()) + if len(versions) == 1: + return {"CUDA Driver Version": versions.pop()} + else: + return {"CUDA Driver Versions": ", ".join(sorted(versions))} + except (OSError, subprocess.SubprocessError): + return {"CUDA Driver Version": "Not Available"} + elif is_rocm_build(): + try: + output = subprocess.check_output( + [ + "rocm-smi", + "--showdriverversion", + "--csv", + ], + text=True, + ) + versions = set(output.strip().splitlines()) + versions.discard("name, value") + if not versions: + return {"ROCM Driver Version": "Not Available"} + ver = versions.pop() + ver = ver.replace('"Driver version", ', "").replace('"', "") + + return {"ROCM Driver Version": ver} + except (OSError, subprocess.SubprocessError): + return {"ROCM Driver Version": "Not Available"} + else: + return {"CUDA Driver Version": "Not Available"} + + +def get_gpu_topology() -> str | None: + """Get GPU topology information.""" + if is_cuda_build(): + try: + result = subprocess.run( + ["nvidia-smi", "topo", "-m"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + check=True, + ) + return "\n" + result.stdout + except (OSError, subprocess.SubprocessError): + return None + elif is_rocm_build(): + try: + result = subprocess.run( + ["rocm-smi", "--showtopotype"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + check=True, + ) + return "\n" + result.stdout + except (OSError, subprocess.SubprocessError): + return None + else: + return None + + +def get_hypervisor_vendor() -> str | None: + try: + output = subprocess.check_output(["lscpu"], text=True) + for line in output.splitlines(): + if "Hypervisor vendor:" in line: + _, _, vendor = line.partition(":") + return vendor.strip() + return None + except (OSError, subprocess.SubprocessError): + return None + + +def main() -> None: + """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_build(): + env_info["NVIDIA Topology"] = gpu_topo + elif is_rocm_build(): + 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__": + main() diff --git a/python/tokenspeed/runtime/cache/__init__.py b/python/tokenspeed/runtime/cache/__init__.py new file mode 100644 index 0000000..11e4a02 --- /dev/null +++ b/python/tokenspeed/runtime/cache/__init__.py @@ -0,0 +1,26 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Runtime cache subsystem. + +This package groups KV cache data structures, allocators, storage backends, +and cache-operation executors under one top-level domain. It also holds the +vision-embedding cache used by the EPD encode stage (:mod:`embedding_cache`). +""" diff --git a/python/tokenspeed/runtime/cache/allocator.py b/python/tokenspeed/runtime/cache/allocator.py new file mode 100755 index 0000000..8ccbf91 --- /dev/null +++ b/python/tokenspeed/runtime/cache/allocator.py @@ -0,0 +1,224 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Allocators for KV-cache page metadata and slot management.""" + +import torch + +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +class KVAllocator: + """Operate on token slots and block-table metadata only. + + Physical KV storage is managed separately so the same metadata operations + can work with different memory backends. + """ + + def __init__( + self, + size: int, + device: str, + max_batch_size: int, + max_context_len: int, + page_size: int, + ): + self.free_slots = None + self.token_slot_refs = None + self.size = size + self.is_not_in_free_group = True + self.free_group = [] + self.page_size = page_size + self.device = device + self.last_slot = torch.ones(max_batch_size, dtype=torch.int32) * (page_size - 1) + self.num_pages = torch.zeros(max_batch_size, dtype=torch.int32) + self.max_context_len = max_context_len + self.max_page_num = (max_context_len + page_size - 1) // page_size + self.max_batch_size = max_batch_size + self.req_to_page = None + self.req_to_page_cpu = None + self.clear() + + def available_size(self) -> int: + return len(self.free_slots) + + def alloc(self, req_pool_index: int, need_size: int, alloced_len: int): + page_offset = alloced_len % self.page_size + page_num = (alloced_len + self.page_size - 1) // self.page_size + last_page_remain = page_num * self.page_size - alloced_len + last_page_id = self.req_to_page_cpu[req_pool_index, page_num - 1].item() + # if last_page_remain is zero, kv_loc is Tensor([]) + kv_loc = ( + last_page_id * self.page_size + + page_offset + + torch.arange(0, min(last_page_remain, need_size), dtype=torch.int32) + ) + if last_page_remain >= need_size: + return kv_loc.to(self.device, non_blocking=True) + + remain_size = need_size - last_page_remain + need_new_page_num = (remain_size + self.page_size - 1) // self.page_size + if need_new_page_num > len(self.free_slots): + # do not change self.seq_lens + return None + + # Check if we have enough space in req_to_page tensor + if page_num + need_new_page_num > self.max_page_num: + logger.warning( + "Requested page range [%s:%s] exceeds max_page_num %s. alloced_len=%s, need_size=%s, page_num=%s", + page_num, + page_num + need_new_page_num, + self.max_page_num, + alloced_len, + need_size, + page_num, + ) + # Do not change self.seq_lens + return None + + new_pages = self.free_slots[:need_new_page_num] + self.free_slots = self.free_slots[need_new_page_num:] + # update req_to_page + self.req_to_page[req_pool_index, page_num : page_num + need_new_page_num] = ( + new_pages.to(self.device) + ) + self.req_to_page_cpu[ + req_pool_index, page_num : page_num + need_new_page_num + ] = new_pages + # construct kv_loc + kv_loc1 = new_pages.unsqueeze(1) * self.page_size + offsets = torch.arange(0, self.page_size, dtype=torch.int32) + kv_loc1 = kv_loc1 + offsets + kv_loc1 = kv_loc1.flatten()[:remain_size] + final_kv_loc = torch.concat([kv_loc, kv_loc1]).to( + self.device, non_blocking=True + ) + return final_kv_loc + + def free_extra_pages_not_cached( + self, req_pool_index: int, real_seq_len: int, alloced_len: int + ): + full_page_num = real_seq_len // self.page_size + alloced_page_num = (alloced_len + self.page_size - 1) // self.page_size + page_num_to_free = alloced_page_num - full_page_num + if page_num_to_free == 0: + return + page_ids_to_free = self.req_to_page[ + req_pool_index, full_page_num : full_page_num + page_num_to_free + ] + self.need_to_free.append(page_ids_to_free) + + def free_req_cache(self, req_pool_index: int, alloced_len: int): + """Release all pages of the request when prefix cache is not used.""" + alloced_page_num = (alloced_len + self.page_size - 1) // self.page_size + if alloced_page_num == 0: + return + page_ids_to_free = self.req_to_page[req_pool_index, :alloced_page_num] + self.need_to_free.append(page_ids_to_free) + + def free_with_diff(self, new_prefix_page_ids, old_page_ids): + # New KV pages come from the prefix tree and are already cached, so only + # release the pages that differ from the old allocation. + if len(new_prefix_page_ids) != len(old_page_ids): + raise ValueError( + "[free with diff] new_prefix_page_ids and old_page_ids " + "should have the same length" + ) + diff = new_prefix_page_ids != old_page_ids + if torch.any(diff): + logger.debug( + "[DebugTrace] free_with_diff free page=%s", old_page_ids[diff].tolist() + ) + self.need_to_free.append(old_page_ids[diff]) + else: + logger.debug( + "[DebugTrace] free_with_diff: no pages to free, all pages are cached" + ) + return diff + + def append_to_later_free(self, page_ids: torch.Tensor) -> None: + self.need_to_free.append(page_ids) + + def free(self, req_pool_index: int, indices=None) -> None: + if self.is_not_in_free_group: + num_pages = self.num_pages[req_pool_index] + pages = self.req_to_page[req_pool_index, :num_pages].cpu() + free_slots = [self.free_slots] + for i in range(num_pages): + page_index = pages[i] + free_slots.append( + torch.arange( + page_index * self.page_size, + (page_index + 1) * self.page_size, + dtype=torch.int32, + ) + ) + self.free_slots = torch.concat(free_slots) + self.num_pages[req_pool_index] = 0 + self.last_slot[req_pool_index] = self.page_size - 1 + else: + self.free_group.append(req_pool_index) + + def free_group_end(self) -> None: + self.is_not_in_free_group = True + if self.need_to_free: + pages_need_to_free = torch.concat(self.need_to_free) + logger.debug( + "[DebugTrace] free_group_end pages_need_to_free=%s", + pages_need_to_free.tolist(), + ) + token_level_offsets = torch.arange(self.page_size, device=self.device) + slots_to_free = ( + pages_need_to_free[:, None] * self.page_size + token_level_offsets + ).flatten() + writted_positions = slots_to_free[self.token_slot_refs[slots_to_free] >= 1] + self.token_slot_refs[writted_positions] += -1 + self.free_slots = torch.concat([self.free_slots, pages_need_to_free.cpu()]) + self.need_to_free = [] + + def clear(self) -> None: + # Page 0 is used for padding + self.free_slots = torch.arange( + 1, self.size // self.page_size, dtype=torch.int32 + ) + if self.token_slot_refs is None: + self.token_slot_refs = torch.zeros( + self.size, dtype=torch.int32, device=self.device + ) + else: + self.token_slot_refs.zero_() + if self.req_to_page is None: + self.req_to_page = torch.zeros( + (self.max_batch_size, self.max_page_num), + dtype=torch.int32, + device=self.device, + ) + self.req_to_page_cpu = torch.zeros( + (self.max_batch_size, self.max_page_num), + dtype=torch.int32, + pin_memory=True, + ) + else: + self.req_to_page.zero_() + self.req_to_page_cpu.zero_() + self.free_group = [] + self.need_to_free = [] diff --git a/python/tokenspeed/runtime/cache/base_prefix_cache.py b/python/tokenspeed/runtime/cache/base_prefix_cache.py new file mode 100755 index 0000000..ada0e94 --- /dev/null +++ b/python/tokenspeed/runtime/cache/base_prefix_cache.py @@ -0,0 +1,107 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Callable +from typing import Any, NamedTuple + +import torch + + +class MatchResult(NamedTuple): + """Result of a prefix match operation. + + Attributes: + device_indices : Page indices of the KV cache on the device matched by common prefix. + last_device_node: The last TreeNode on the device that was matched. + device_prefix_length: Length of the common prefix in tokens (not pages). + last_host_node : The last TreeNode on the host that was matched. + Note that if KVStore is not enabled, + this **must** be the same as `last_device_node`. + host_hit_length : Number of tokens hit on the host, if applicable. + 0 if KVStore is not enabled. + Note: node.host_value stores token indices. + """ + + device_indices: torch.Tensor | None = None + last_device_node: Any = None + device_prefix_length: int = 0 + last_host_node: Any = None + host_hit_length: int = 0 + + +class BasePrefixCache(ABC): + """Cache can be indexed by either rid or key.""" + + @abstractmethod + def reset(self) -> None: + raise NotImplementedError + + @abstractmethod + def match_prefix(self, **kwargs) -> MatchResult: + """Match a request prefix and optionally prepare req-local state. + + Unified contract: + - When called without `req`, implementations should behave like a pure + prefix lookup and avoid mutating request-local state. + - When called with `req`, implementations may prepare request-local + execution state required by that cache backend. + """ + raise NotImplementedError + + @abstractmethod + def insert(self, **kwargs): + raise NotImplementedError + + @abstractmethod + def cache_finished_req(self, **kwargs): + raise NotImplementedError + + @abstractmethod + def cache_unfinished_req(self, **kwargs): + raise NotImplementedError + + @abstractmethod + def evict(self, num_tokens: int, evict_callback: Callable[..., None]): + raise NotImplementedError + + @abstractmethod + def inc_lock_ref(self, node): + raise NotImplementedError + + @abstractmethod + def dec_lock_ref(self, node): + raise NotImplementedError + + @abstractmethod + def evictable_size(self): + raise NotImplementedError + + @abstractmethod + def protected_size(self): + raise NotImplementedError() + + def total_size(self): + raise NotImplementedError() + + def pretty_print(self): + raise NotImplementedError() diff --git a/python/tokenspeed/runtime/cache/embedding_cache.py b/python/tokenspeed/runtime/cache/embedding_cache.py new file mode 100644 index 0000000..d649255 --- /dev/null +++ b/python/tokenspeed/runtime/cache/embedding_cache.py @@ -0,0 +1,224 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Vision-embedding cache for the EPD encode stage. + +The encode server runs the vision tower only; a cache lets duplicate images +(same ``MultimodalDataItem.hash``) reuse a previously computed embedding instead +of re-running the tower. + +Two interchangeable implementations share a ``get`` / ``put`` surface: +:class:`EmbeddingCache` (single-tier, bytes-bounded LRU) and +:class:`TieredEmbeddingCache` (L1 GPU VRAM + L2 host DRAM). Both are +framework-agnostic and unit-testable without a GPU -- the caller supplies byte +sizes and the device<->host copies are injectable. +""" + +from __future__ import annotations + +import collections +from collections.abc import Callable, Hashable + + +class EmbeddingCache: + """Bytes-bounded LRU cache of vision-tower outputs, keyed by content hash. + + Values are opaque (typically a ``torch.Tensor``); the caller supplies the + byte size so this stays framework-agnostic and testable without a GPU. + + ``on_evict`` is an optional ``(key, value, nbytes)`` callback fired only for + an entry dropped by capacity overflow (not an update to an existing key, + which is an in-place replace, nor an explicit :meth:`pop`). It lets a second + tier capture LRU victims and demote them (see :class:`TieredEmbeddingCache`). + """ + + def __init__( + self, + capacity_bytes: int, + on_evict: Callable[[Hashable, object, int], None] | None = None, + ): + if capacity_bytes < 0: + raise ValueError(f"capacity_bytes must be >= 0, got {capacity_bytes}") + self.capacity_bytes = capacity_bytes + self._on_evict = on_evict + # key -> (value, nbytes); ordered by recency (oldest first). + self._store: collections.OrderedDict[Hashable, tuple[object, int]] = ( + collections.OrderedDict() + ) + self._cur_bytes = 0 + self.hits = 0 + self.misses = 0 + + def get(self, key: Hashable) -> object | None: + entry = self._store.get(key) + if entry is None: + self.misses += 1 + return None + self._store.move_to_end(key) + self.hits += 1 + return entry[0] + + def put(self, key: Hashable, value: object, nbytes: int) -> None: + if nbytes < 0: + raise ValueError(f"nbytes must be >= 0, got {nbytes}") + # An item larger than the whole cache can never be retained; skip it + # rather than evicting everything for a value dropped immediately after. + if nbytes > self.capacity_bytes: + return + existing = self._store.pop(key, None) + if existing is not None: + self._cur_bytes -= existing[1] + self._store[key] = (value, nbytes) + self._cur_bytes += nbytes + self._evict() + + def pop(self, key: Hashable) -> tuple[object, int] | None: + """Remove ``key`` and return its ``(value, nbytes)``, or ``None`` if + absent. A structural removal: it neither counts as a hit/miss nor fires + ``on_evict`` (the caller is taking ownership of the value, e.g. to + promote it to another tier).""" + existing = self._store.pop(key, None) + if existing is None: + return None + self._cur_bytes -= existing[1] + return existing + + def _evict(self) -> None: + while self._cur_bytes > self.capacity_bytes and self._store: + key, (value, nbytes) = self._store.popitem(last=False) + self._cur_bytes -= nbytes + if self._on_evict is not None: + self._on_evict(key, value, nbytes) + + def __contains__(self, key: Hashable) -> bool: + return key in self._store + + def __len__(self) -> int: + return len(self._store) + + @property + def current_bytes(self) -> int: + return self._cur_bytes + + +def _embedding_to_host(value: object) -> object: + """Default L1->L2 demote: copy a vision embedding from GPU to host memory. + + A cache value is a ``(main, deepstack)`` tuple (deepstack half may be + ``None`` for non-deepstack models) or a bare tensor. + """ + if isinstance(value, tuple): + return tuple(None if t is None else t.cpu() for t in value) + return value.cpu() + + +def _embedding_to_device(value: object, device) -> object: + """Default L2->L1 promote: copy a host-resident embedding back to ``device`` + so the executor can stage it into the GPU send ring.""" + if isinstance(value, tuple): + return tuple(None if t is None else t.to(device) for t in value) + return value.to(device) + + +class TieredEmbeddingCache: + """Two-tier vision-embedding cache: L1 in GPU VRAM, L2 in host DRAM. + + Exposes the same ``get`` / ``put`` surface as :class:`EmbeddingCache`, so the + encode worker uses either interchangeably. L2 catches L1's LRU victims in + cheaper host DRAM: a hit there skips the (much more expensive) ViT and only + pays a host->device copy. + + Tiers are kept *exclusive* -- a key lives in exactly one. An L1 eviction + demotes the victim to L2 (device->host copy); an L2 hit promotes the entry + back to L1 (host->device copy) and removes it from L2; ``put`` always lands + in L1 and drops any stale L2 duplicate. No distributed L3: image-hash routing + pins each image to one worker, so a per-worker local cache captures the reuse. + + The device<->host copies are injectable (``to_host`` / ``to_device``) so the + tiering logic is unit-testable without a GPU; the defaults copy real tensors + to/from ``device``. + """ + + def __init__( + self, + l1_capacity_bytes: int, + l2_capacity_bytes: int, + *, + device: object = None, + to_host: Callable[[object], object] | None = None, + to_device: Callable[[object], object] | None = None, + ): + # L1 demotes its evictions into L2 via the on_evict hook; L2 is the + # bottom tier (no hook), so its evictions are true drops -- no recursion. + self.l1 = EmbeddingCache(l1_capacity_bytes, on_evict=self._demote) + self.l2 = EmbeddingCache(l2_capacity_bytes) + self._device = device + self._to_host = to_host or _embedding_to_host + self._to_device = to_device or (lambda v: _embedding_to_device(v, self._device)) + self.l1_hits = 0 + self.l2_hits = 0 + self.misses = 0 + self.promotions = 0 + self.demotions = 0 + + def get(self, key: Hashable) -> object | None: + value = self.l1.get(key) + if value is not None: + self.l1_hits += 1 + return value + promoted = self.l2.pop(key) + if promoted is None: + self.misses += 1 + return None + host_value, nbytes = promoted + device_value = self._to_device(host_value) + self.l2_hits += 1 + self.promotions += 1 + # Re-home as MRU in L1. May evict colder L1 entries, which demote to L2; + # the just-promoted key is MRU so it is never the victim. + self.l1.put(key, device_value, nbytes) + return device_value + + def put(self, key: Hashable, value: object, nbytes: int) -> None: + # L1 is the write tier. Drop any stale L2 copy first so the tiers stay + # exclusive (e.g. a key demoted earlier and now re-encoded on a miss). + self.l2.pop(key) + self.l1.put(key, value, nbytes) + + def _demote(self, key: Hashable, value: object, nbytes: int) -> None: + """L1 eviction hook: stash the victim in host DRAM instead of dropping + it. host nbytes == device nbytes (same dtype/numel). A victim larger than + all of L2 is silently dropped by ``L2.put`` (same as having no L2).""" + # Skip the wasted device->host copy for a victim L2 would just drop. + if nbytes > self.l2.capacity_bytes: + return + self.l2.put(key, self._to_host(value), nbytes) + if key in self.l2: + self.demotions += 1 + + @property + def hits(self) -> int: + return self.l1_hits + self.l2_hits + + def __contains__(self, key: Hashable) -> bool: + return key in self.l1 or key in self.l2 + + def __len__(self) -> int: + return len(self.l1) + len(self.l2) diff --git a/python/tokenspeed/runtime/cache/evict_policy.py b/python/tokenspeed/runtime/cache/evict_policy.py new file mode 100755 index 0000000..915e1d2 --- /dev/null +++ b/python/tokenspeed/runtime/cache/evict_policy.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the FluentLLM project +# SPDX-FileCopyrightText: Copyright 2023-2024 SGLang Team +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Eviction strategy definitions for cache tree nodes.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from tokenspeed.runtime.cache.prefix_cache import TreeNode + + +class EvictionStrategy(ABC): + @abstractmethod + def get_priority(self, node: "TreeNode") -> float | tuple[Any, ...]: + """Return the sortable priority used for eviction.""" + + +class LRUStrategy(EvictionStrategy): + def get_priority(self, node: "TreeNode") -> float: + return node.last_access_time + + +class LFUStrategy(EvictionStrategy): + def get_priority(self, node: "TreeNode") -> tuple[int, float]: + return (node.hit_count, node.last_access_time) + + +class FIFOStrategy(EvictionStrategy): + def get_priority(self, node: "TreeNode") -> float: + return node.creation_time + + +class MRUStrategy(EvictionStrategy): + def get_priority(self, node: "TreeNode") -> float: + return -node.last_access_time + + +class FILOStrategy(EvictionStrategy): + def get_priority(self, node: "TreeNode") -> float: + return -node.creation_time + + +class PriorityStrategy(EvictionStrategy): + """Priority-aware eviction with LRU tiebreaking.""" + + def get_priority(self, node: "TreeNode") -> tuple[int, float]: + # Lower priority values are evicted first; ties fall back to LRU. + return (node.priority, node.last_access_time) diff --git a/python/tokenspeed/runtime/cache/executor/__init__.py b/python/tokenspeed/runtime/cache/executor/__init__.py new file mode 100644 index 0000000..60f971d --- /dev/null +++ b/python/tokenspeed/runtime/cache/executor/__init__.py @@ -0,0 +1,21 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Cache-operation executors for host and storage transfers.""" diff --git a/python/tokenspeed/runtime/cache/executor/flat_memory_executor.py b/python/tokenspeed/runtime/cache/executor/flat_memory_executor.py new file mode 100644 index 0000000..b9e100d --- /dev/null +++ b/python/tokenspeed/runtime/cache/executor/flat_memory_executor.py @@ -0,0 +1,445 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Flat host-tier executor (M15 Phase D2): drives FlatWriteBack/FlatLoadBack +page-id pairs against the byte-blind :class:`FlatHostMirror`, replacing the +radix ``MemoryExecutor`` when serving with a flat-built scheduler ext and the +kvstore enabled. Unlike the radix host executor it ACKS loadbacks: the flat +C++ scheduler pins source host pages and destination device blocks until a +``Cache.LoadBackDoneEvent`` retires the op. +""" + +from __future__ import annotations + +from collections import OrderedDict +from collections.abc import Iterable, Sequence + +import psutil +import torch +from tokenspeed_scheduler import Cache + +from tokenspeed.runtime.cache.executor.host_executor import ( + _Ack, + _cache_stream_priorities, + _new_cache_stream, + _ordered_unique, +) +from tokenspeed.runtime.cache.flat_host_mirror import ( + FlatHostMirror, + flat_bytes_per_host_page, +) +from tokenspeed.runtime.cache.kvstore_controller import LayerDoneCounter +from tokenspeed.runtime.cache.transfer.types import CacheKind +from tokenspeed.runtime.execution.cuda_graph_wrapper import get_is_capture_mode +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + +_HOST_MEM_HEADROOM_BYTES = 10 * (1024**3) + + +def flat_num_host_pages( + *, + bytes_per_host_page: int, + device_pool_size: int, + page_size: int, + host_ratio: float, + host_size_gb: float, +) -> int: + """Host page budget from the kvstore sizing knobs (same knobs the radix + ``HostKVCache`` resolves, kv_cache_host.py:91-102, budget arithmetic only): + + - ``host_size_gb > 0``: explicit byte budget, floor to whole mirror pages + (never exceeds the requested bytes): + ``host_size_gb * 1e9 // bytes_per_host_page``. + - otherwise ratio sizing, mirroring the radix token->page align-up: + ``int(device_pool_size * host_ratio) // page_size + 1``. + """ + if bytes_per_host_page <= 0: + raise ValueError(f"bytes_per_host_page must be > 0, got {bytes_per_host_page}") + if page_size <= 0: + raise ValueError(f"page_size must be > 0, got {page_size}") + if host_size_gb > 0: + num_pages = int(host_size_gb * 1e9 // bytes_per_host_page) + else: + num_pages = int(device_pool_size * host_ratio) // page_size + 1 + if num_pages <= 0: + raise ValueError( + "flat host tier resolved to zero host pages " + f"(host_size_gb={host_size_gb}, host_ratio={host_ratio}, " + f"bytes_per_host_page={bytes_per_host_page}); increase the " + "kvstore size." + ) + return num_pages + + +class FlatMemoryExecutor: + """Slim replacement for ``MemoryExecutor`` under the flat host tier. + + Exposes the exact surface ``EventLoop`` drives: ``submit_plan`` / + ``poll_results`` / ``get_producer_index`` / ``set_consumer`` (plus the + ``host_exec.pools`` attribute walk in ``_setup_layerwise_loadback``). + No host pool, no storage executor, no mamba: the flat scheduler config + validation already rejects those setups. + """ + + # EventLoop keys per-op inflight accounting off this: flat loadbacks are + # acked (LoadBackDoneEvent), radix loadbacks are not. + emits_loadback_acks = True + + def __init__(self, device_pool, *, host_ratio: float, host_size_gb: float): + self.page_size = int(device_pool.page_size) + self.layer_num = len(device_pool.k_buffer) + + bytes_per_host_page = flat_bytes_per_host_page(device_pool) + num_host_pages = flat_num_host_pages( + bytes_per_host_page=bytes_per_host_page, + device_pool_size=int(device_pool.size), + page_size=self.page_size, + host_ratio=host_ratio, + host_size_gb=host_size_gb, + ) + requested_bytes = num_host_pages * bytes_per_host_page + available_bytes = psutil.virtual_memory().available - _HOST_MEM_HEADROOM_BYTES + if requested_bytes > available_bytes: + raise ValueError( + f"Not enough host memory for the flat host tier. Requesting " + f"{requested_bytes / 1e9:.2f} GB but only have " + f"{available_bytes / 1e9:.2f} GB free. Please reduce the " + f"size of the KVStore." + ) + logger.info( + "Allocating %.2f GB pinned host memory for the flat host tier " + "(num_host_pages=%s bytes_per_host_page=%s host_size_gb=%r " + "host_ratio=%r device_pool.size=%r)", + requested_bytes / 1e9, + num_host_pages, + bytes_per_host_page, + host_size_gb, + host_ratio, + device_pool.size, + ) + self.mirror = FlatHostMirror(device_pool, num_host_pages) + self.num_host_pages = num_host_pages + + # Layerwise loadback fencing: register the counter where the radix + # KVCachePool would, so pool.get_key_buffer/get_value_buffer gate on + # the same wait_until(layer_id) machinery. + self._counter = LayerDoneCounter(self.layer_num) + device_pool.register_layer_transfer_counter(self._counter) + # _start_loading maps layer -> mirror V-tensor event and relies on + # load_events[-1] (LayerLoadingEvent.finish_event, reuse fence in + # update_producer) covering EVERY copy: it pins the last layer to + # the op's last per-tensor event, which without state slabs is the + # last layer's V event only if that V mirror is the last KV tensor + # pair. Holds for both layouts (legacy: identity; slab: last layer + # is the last occurrence of its group). A state last layer has no + # KV mirror at all -- its copies (state slabs trail every KV + # tensor) are covered by the events[-1] pin in _start_loading. + assert ( + self.mirror.state_tensor_indices_of_layer(self.layer_num - 1) is not None + or self.mirror.tensor_index_of_layer(self.layer_num - 1) + == self.mirror.num_k_tensors - 1 + ), "flat host tier: last layer's V mirror is not the last KV tensor pair" + + write_priority, load_priority = _cache_stream_priorities() + self.write_stream = _new_cache_stream(write_priority) + self.load_stream = _new_cache_stream(load_priority) + + # (device_page, host_page) pairs staged between submit() and flush(). + self._pending_write_pairs: list[tuple[int, int]] = [] + self._pending_write_op_ids: list[int] = [] + self._pending_load_pairs: list[tuple[int, int]] = [] + self._pending_load_op_ids: list[int] = [] + self.ack_write_queue: list[_Ack] = [] + self.ack_load_queue: list[_Ack] = [] + # Ops whose page lists were empty on the wire (C++ dedups transfers + # across ops of one batched operation) and no batch event covers them. + self._immediate_write_op_ids: list[int] = [] + self._immediate_load_op_ids: list[int] = [] + + self._producer_map: OrderedDict[int, int] = OrderedDict() + self._producer_map_limit = 1024 + + # Surface for EventLoop._setup_layerwise_loadback, which walks + # memory_executor.host_exec.pools to enumerate fencing kinds. + self.host_exec = self + self.pools = {CacheKind.KV: self.mirror} + + # ------------------------------------------------------------------ + # Submission (wire shape: batched Flat{WriteBack,LoadBack}Operation) + # ------------------------------------------------------------------ + + def submit_plan(self, plan) -> None: + if plan.cache: + logger.debug("[cache_op] flat submit_plan: %s cache ops", len(plan.cache)) + for op in plan.cache: + self.submit(op) + self.flush() + + def submit(self, op) -> None: + if isinstance(op, Cache.WriteBackOp): + self.submit_writeback(op.op_ids, op.src_pages, op.dst_pages) + elif isinstance(op, Cache.LoadBackOp): + self.submit_loadback(op.op_ids, op.src_pages, op.dst_pages) + else: + raise ValueError( + f"flat host tier: unsupported cache op kind {type(op).__name__}" + ) + + def _submit( + self, + op_ids: Sequence[int], + src_pages: Sequence[Sequence[int]], + dst_pages: Sequence[Sequence[int]], + *, + pending_op_ids: list[int], + pending_pairs: list[tuple[int, int]], + src_is_device: bool, + ) -> None: + """Stage copies as (device_page, host_page) pairs; fail loud on a + ragged wire payload instead of silently dropping trailing ops.""" + assert len(op_ids) == len(src_pages) == len(dst_pages), ( + f"flat host tier: ragged cache-op payload (op_ids={len(op_ids)}, " + f"src_pages={len(src_pages)}, dst_pages={len(dst_pages)})" + ) + for op_id, src, dst in zip(op_ids, src_pages, dst_pages): + assert len(src) == len(dst), ( + f"flat host tier: op {op_id} src/dst page lists differ " + f"({len(src)} vs {len(dst)})" + ) + pending_op_ids.append(int(op_id)) + device_pages, host_pages = (src, dst) if src_is_device else (dst, src) + pending_pairs.extend( + (int(d), int(h)) for d, h in zip(device_pages, host_pages) + ) + + def submit_writeback( + self, + op_ids: Sequence[int], + src_pages: Sequence[Sequence[int]], + dst_pages: Sequence[Sequence[int]], + ) -> None: + """Stage device->host copies: src=device pages, dst=host pages.""" + self._submit( + op_ids, + src_pages, + dst_pages, + pending_op_ids=self._pending_write_op_ids, + pending_pairs=self._pending_write_pairs, + src_is_device=True, + ) + + def submit_loadback( + self, + op_ids: Sequence[int], + src_pages: Sequence[Sequence[int]], + dst_pages: Sequence[Sequence[int]], + ) -> None: + """Stage host->device copies: src=host pages, dst=device pages.""" + self._submit( + op_ids, + src_pages, + dst_pages, + pending_op_ids=self._pending_load_op_ids, + pending_pairs=self._pending_load_pairs, + src_is_device=False, + ) + + def flush(self) -> None: + self._start_loading() + self._start_writing() + + def _start_writing(self) -> None: + if not self._pending_write_op_ids: + return + op_ids = _ordered_unique(self._pending_write_op_ids) + pairs = self._pending_write_pairs + self._pending_write_op_ids = [] + self._pending_write_pairs = [] + if not pairs: + self._immediate_write_op_ids.extend(op_ids) + return + # Order the D2H copies after already-enqueued default-stream work + # (same fence the radix _start_writing places). + start_event = torch.cuda.Event() + start_event.record() + start_event.wait(self.write_stream) + self.mirror.store_pages(pairs, self.write_stream) + finish_event = torch.cuda.Event() + finish_event.record(self.write_stream) + self.ack_write_queue.append(_Ack(finish_event, op_ids)) + + def _start_loading(self) -> None: + if not self._pending_load_op_ids: + return + assert ( + not get_is_capture_mode() + ), "cache loadback must run in eager admission iter" + op_ids = _ordered_unique(self._pending_load_op_ids) + pairs = self._pending_load_pairs + self._pending_load_op_ids = [] + self._pending_load_pairs = [] + if not pairs: + self._immediate_load_op_ids.extend(op_ids) + return + + producer_id = self._counter.update_producer() + producer_event = self._counter.events[producer_id] + producer_event.start_event.record() + producer_event.start_event.wait(self.load_stream) + + events = self.mirror.load_pages_with_events(pairs, self.load_stream) + # Layer fence: layer L is readable once its V mirror copy lands; the + # load stream is serial (all K copies precede all V copies), so the + # V-tensor event also covers L's K copy. Paired slab layers share the + # slab event -- correct by design. State layers instead fence on + # their ssm event: conv precedes ssm in tensor_pairs order (and both + # follow every KV tensor), so on the serial stream the ssm event + # covers the conv copy and the layer's KV copies -- the same + # K-before-V reasoning as above. + num_k = self.mirror.num_k_tensors + for layer_id in range(self.layer_num): + state_indices = self.mirror.state_tensor_indices_of_layer(layer_id) + if state_indices is not None: + producer_event.load_events[layer_id] = events[state_indices[1]] + else: + producer_event.load_events[layer_id] = events[ + num_k + self.mirror.tensor_index_of_layer(layer_id) + ] + # finish_event (== load_events[-1]) is the producer-slot reuse fence + # in update_producer and must cover EVERY copy of the op; state + # tensors copy after all KV tensors, so pin the last layer to the + # op's last per-tensor event. A no-op without state slabs: events[-1] + # is then the last layer's V event (ctor assert). + producer_event.load_events[self.layer_num - 1] = events[-1] + # events[-1] is also the reassigned finish_event, so the ack covers + # every copy. + self.ack_load_queue.append(_Ack(events[-1], op_ids)) + for op_id in op_ids: + self._producer_map[op_id] = producer_id + while len(self._producer_map) > self._producer_map_limit: + self._producer_map.popitem(last=False) + + # ------------------------------------------------------------------ + # Ack draining + # ------------------------------------------------------------------ + + def poll_results(self) -> list: + results: list = [] + for op_id in self._immediate_write_op_ids: + results.append(self._write_done(op_id)) + self._immediate_write_op_ids.clear() + for op_id in self._immediate_load_op_ids: + results.append(self._load_done(op_id)) + self._immediate_load_op_ids.clear() + + remaining_writes = [] + for ack in self.ack_write_queue: + if ack.finish_event.query(): + results.extend(self._write_done(op_id) for op_id in ack.op_ids) + else: + remaining_writes.append(ack) + self.ack_write_queue[:] = remaining_writes + + remaining_loads = [] + for ack in self.ack_load_queue: + if ack.finish_event.query(): + results.extend(self._load_done(op_id) for op_id in ack.op_ids) + else: + remaining_loads.append(ack) + self.ack_load_queue[:] = remaining_loads + + if results: + for r in results: + logger.debug( + "[cache_op] flat done op_id=%s success=%s type=%s", + r.op_id, + r.success, + type(r).__name__, + ) + return results + + @staticmethod + def _write_done(op_id: int): + evt = Cache.WriteBackDoneEvent() + evt.op_id = op_id + evt.success = True + return evt + + @staticmethod + def _load_done(op_id: int): + evt = Cache.LoadBackDoneEvent() + evt.op_id = op_id + evt.success = True + return evt + + # ------------------------------------------------------------------ + # Layerwise loadback fencing (EventLoop._setup_layerwise_loadback) + # ------------------------------------------------------------------ + + def get_producer_index( + self, kind_or_op_id: CacheKind | str | int, op_id: int | None = None + ) -> int | None: + if op_id is None: + op_id = int(kind_or_op_id) + return self._producer_map.pop(int(op_id), None) + + def set_consumer( + self, + kind_or_producer_index: CacheKind | str | int | Iterable[int], + producer_index: int | Iterable[int] | None = None, + ) -> None: + if producer_index is None: + producer_index = kind_or_producer_index + self._counter.set_consumer(producer_index) + + # ------------------------------------------------------------------ + # MemoryExecutor surface stubs + # ------------------------------------------------------------------ + + def set_mamba_layerwise_cow(self, cow_dst_pages_by_src) -> None: + assert not cow_dst_pages_by_src, ( + "flat host tier has no mamba L2: the flat scheduler config " + "validation rejects state-only pools" + ) + + def query_l3_pages(self, hashes: list[str]) -> int: + # No L3 storage tier under the flat build (EventLoop refuses a + # storage backend up front); report zero hits. + return 0 + + def shutdown(self) -> None: + self.write_stream.synchronize() + self.load_stream.synchronize() + + def reset(self) -> None: + self.write_stream.synchronize() + self.load_stream.synchronize() + self._pending_write_pairs.clear() + self._pending_write_op_ids.clear() + self._pending_load_pairs.clear() + self._pending_load_op_ids.clear() + self.ack_write_queue.clear() + self.ack_load_queue.clear() + self._immediate_write_op_ids.clear() + self._immediate_load_op_ids.clear() + self._producer_map.clear() + self._counter.reset() diff --git a/python/tokenspeed/runtime/cache/executor/host_executor.py b/python/tokenspeed/runtime/cache/executor/host_executor.py new file mode 100644 index 0000000..2830daf --- /dev/null +++ b/python/tokenspeed/runtime/cache/executor/host_executor.py @@ -0,0 +1,516 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Host-side executor for cache writeback and loadback operations.""" + +from __future__ import annotations + +from collections import OrderedDict +from collections.abc import Iterable +from typing import NamedTuple + +import torch +from tokenspeed_scheduler import Cache + +from tokenspeed.runtime.cache.transfer.kv_pool import KVCachePool +from tokenspeed.runtime.cache.transfer.pool import CachePool +from tokenspeed.runtime.cache.transfer.types import CacheKind, Location, TransferUnit +from tokenspeed.runtime.execution.cuda_graph_wrapper import get_is_capture_mode +from tokenspeed.runtime.utils import get_colorful_logger, get_device_module + +logger = get_colorful_logger(__name__) +device_module = get_device_module() +CONCURRENT_WRITEBACK_BLOCK_QUOTA = 2 + + +def _cache_stream_priorities() -> tuple[int | None, int | None]: + priority_range = getattr(device_module.Stream, "priority_range", None) + if priority_range is None: + return None, None + try: + least_priority, greatest_priority = priority_range() + except (RuntimeError, TypeError): + return None, None + return least_priority, greatest_priority + + +def _new_cache_stream(priority: int | None = None): + if priority is None: + return device_module.Stream() + try: + return device_module.Stream(priority=priority) + except (RuntimeError, TypeError): + return device_module.Stream() + + +def page_ids_to_token_indices( + page_ids: list[int], + page_size: int, + device: str = "cpu", +) -> torch.Tensor: + if len(page_ids) == 0: + return torch.empty((0,), dtype=torch.int64, device=device) + pages = torch.tensor(page_ids, dtype=torch.int64, device=device) + offsets = torch.arange(page_size, dtype=torch.int64, device=device) + return (pages[:, None] * page_size + offsets[None, :]).reshape(-1) + + +def _dedupe_page_pairs( + src_pages: Iterable[int], + dst_pages: Iterable[int], +) -> tuple[list[int], list[int]]: + seen = set() + deduped_src = [] + deduped_dst = [] + for src_page, dst_page in zip(src_pages, dst_pages): + pair = (int(src_page), int(dst_page)) + if pair in seen: + continue + seen.add(pair) + deduped_src.append(pair[0]) + deduped_dst.append(pair[1]) + return deduped_src, deduped_dst + + +def _ordered_unique(values: Iterable[int]) -> list[int]: + seen = set() + result = [] + for value in values: + value = int(value) + if value in seen: + continue + seen.add(value) + result.append(value) + return result + + +class _Ack(NamedTuple): + finish_event: object # device_module.Event + op_ids: list[int] + + +class HostExecutor: + def __init__( + self, + page_size: int | None = None, + device_pool=None, + host_pool=None, + io_backend: str = "kernel", + layer_num: int | None = None, + draft_device_pool=None, + draft_host_pool=None, + draft_layer_num: int = 0, + pools: list[CachePool] | None = None, + ): + self.io_backend = io_backend + if pools is None: + if ( + page_size is None + or device_pool is None + or host_pool is None + or layer_num is None + ): + raise ValueError("HostExecutor requires either pools or KV pool inputs") + pools = [ + KVCachePool( + device_pool=device_pool, + host_pool=host_pool, + io_backend=io_backend, + layer_num=layer_num, + draft_device_pool=draft_device_pool, + draft_host_pool=draft_host_pool, + draft_layer_num=draft_layer_num, + ) + ] + if not pools: + raise ValueError("HostExecutor requires at least one cache pool") + + self.pools = {CacheKind(pool.kind): pool for pool in pools} + self.device = next(iter(self.pools.values())).device + + write_priority, load_priority = _cache_stream_priorities() + self.write_stream = _new_cache_stream(write_priority) + self.load_stream = _new_cache_stream(load_priority) + self._writeback_block_quota: int | None = None + + self.write_queues: dict[CacheKind, list[TransferUnit]] = { + kind: [] for kind in self.pools + } + self.load_queues: dict[CacheKind, list[TransferUnit]] = { + kind: [] for kind in self.pools + } + + self.ack_write_queue: list[_Ack] = [] + self.ack_load_queue: list[_Ack] = [] + self.completed_writebacks: list[int] = [] + + self._counters = { + kind: pool.get_layer_done_counter() for kind, pool in self.pools.items() + } + self._producer_map: dict[CacheKind, OrderedDict[int, int]] = { + kind: OrderedDict() for kind in self.pools + } + self._producer_map_limit = 1024 + + def enqueue_writeback( + self, + op_id, + src_pages, + dst_pages, + is_retract: bool = False, + kind: CacheKind | str = CacheKind.KV, + ) -> None: + kind = CacheKind(kind) + pool = self.pools[kind] + src_pages, dst_pages = _dedupe_page_pairs(src_pages, dst_pages) + if not src_pages: + self.completed_writebacks.append(op_id) + return + device_indices = page_ids_to_token_indices(src_pages, pool.page_size(), "cpu") + host_indices = page_ids_to_token_indices(dst_pages, pool.page_size(), "cpu") + self.write_queues[kind].append( + TransferUnit( + kind=kind, + src_loc=Location.DEVICE, + dst_loc=Location.HOST, + src_indices=device_indices, + dst_indices=host_indices, + op_id=op_id, + is_retract=is_retract, + ) + ) + + def enqueue_loadback( + self, + op_id, + src_pages, + dst_pages, + kind: CacheKind | str = CacheKind.KV, + layerwise_cow_dst_pages_by_src: dict[int, list[int]] | None = None, + ) -> None: + kind = CacheKind(kind) + pool = self.pools[kind] + src_pages, dst_pages = _dedupe_page_pairs(src_pages, dst_pages) + if not src_pages: + return + host_indices = page_ids_to_token_indices(src_pages, pool.page_size(), "cpu") + device_indices = page_ids_to_token_indices(dst_pages, pool.page_size(), "cpu") + cow_src_indices = None + cow_dst_indices = None + if layerwise_cow_dst_pages_by_src: + cow_src_pages: list[int] = [] + cow_dst_pages: list[int] = [] + for dst_page in dst_pages: + for cow_dst in layerwise_cow_dst_pages_by_src.get(int(dst_page), []): + cow_src_pages.append(int(dst_page)) + cow_dst_pages.append(int(cow_dst)) + if cow_src_pages: + cow_src_indices = page_ids_to_token_indices( + cow_src_pages, pool.page_size(), "cpu" + ) + cow_dst_indices = page_ids_to_token_indices( + cow_dst_pages, pool.page_size(), "cpu" + ) + self.load_queues[kind].append( + TransferUnit( + kind=kind, + src_loc=Location.HOST, + dst_loc=Location.DEVICE, + src_indices=host_indices, + dst_indices=device_indices, + op_id=op_id, + layerwise_cow_src_indices=cow_src_indices, + layerwise_cow_dst_indices=cow_dst_indices, + ) + ) + + def flush(self) -> None: + throttle_writeback = self._has_work(self.load_queues) and not any( + unit.is_retract for units in self.write_queues.values() for unit in units + ) + writeback_block_quota = ( + CONCURRENT_WRITEBACK_BLOCK_QUOTA if throttle_writeback else None + ) + previous_writeback_block_quota = getattr(self, "_writeback_block_quota", None) + self._writeback_block_quota = writeback_block_quota + try: + self._start_loading() + self._start_writing() + finally: + self._writeback_block_quota = previous_writeback_block_quota + + def _start_writing(self) -> None: + if not self._has_work(self.write_queues): + return + + start_event = device_module.Event() + finish_event = device_module.Event() + op_ids: list[int] = [] + + start_event.record() + with device_module.stream(self.write_stream): + start_event.wait(self.write_stream) + for kind, units in self.write_queues.items(): + if not units: + continue + pool = self.pools[kind] + unit = self._merge_units(units) + src_indices, dst_indices = self._prepare_indices(unit, pool) + self._pool_writeback( + pool, src_indices.to(torch.int64), dst_indices.to(torch.int64) + ) + self._record_if_cuda(src_indices, self.write_stream) + self._record_if_cuda(dst_indices, self.write_stream) + op_ids.extend(unit.op_id for unit in units) + finish_event.record() + + self._clear_queues(self.write_queues) + self.ack_write_queue.append(_Ack(finish_event, _ordered_unique(op_ids))) + + def _start_loading(self) -> None: + if not self._has_work(self.load_queues): + return + assert ( + not get_is_capture_mode() + ), "cache loadback must run in eager admission iter" + + with device_module.stream(self.load_stream): + for kind, units in self.load_queues.items(): + if not units: + continue + pool = self.pools[kind] + counter = self._counters[kind] + producer_id = counter.update_producer() + producer_event = counter.events[producer_id] + producer_event.start_event.record() + producer_event.start_event.wait(self.load_stream) + + unit = self._merge_units(units) + src_indices, dst_indices = self._prepare_indices(unit, pool) + layerwise_copy = getattr(pool, "copy_layer", None) + cow_src_indices = unit.layerwise_cow_src_indices + cow_dst_indices = unit.layerwise_cow_dst_indices + for layer_index in range(pool.num_layers()): + pool.loadback( + src_indices.to(torch.int64), + dst_indices.to(torch.int64), + layer_index, + ) + if ( + layerwise_copy is not None + and cow_src_indices is not None + and cow_dst_indices is not None + ): + layerwise_copy( + cow_src_indices.to(torch.int64), + cow_dst_indices.to(torch.int64), + layer_index, + ) + producer_event.complete(layer_index) + self._record_if_cuda(src_indices, self.load_stream) + self._record_if_cuda(dst_indices, self.load_stream) + if cow_src_indices is not None: + self._record_if_cuda(cow_src_indices, self.load_stream) + if cow_dst_indices is not None: + self._record_if_cuda(cow_dst_indices, self.load_stream) + + op_ids = _ordered_unique(unit.op_id for unit in units) + self.ack_load_queue.append(_Ack(producer_event.finish_event, op_ids)) + producer_map = self._producer_map[kind] + for op_id in op_ids: + producer_map[op_id] = producer_id + while len(producer_map) > self._producer_map_limit: + producer_map.popitem(last=False) + + self._clear_queues(self.load_queues) + + @staticmethod + def _has_work(queues: dict[CacheKind, list[TransferUnit]]) -> bool: + return any(bool(units) for units in queues.values()) + + @staticmethod + def _clear_queues(queues: dict[CacheKind, list[TransferUnit]]) -> None: + for units in queues.values(): + units.clear() + + @staticmethod + def _merge_units(units: list[TransferUnit]) -> TransferUnit: + assert units + if len(units) == 1: + return units[0] + first = units[0] + cow_src_indices = [ + unit.layerwise_cow_src_indices + for unit in units + if unit.layerwise_cow_src_indices is not None + ] + cow_dst_indices = [ + unit.layerwise_cow_dst_indices + for unit in units + if unit.layerwise_cow_dst_indices is not None + ] + return TransferUnit( + kind=first.kind, + src_loc=first.src_loc, + dst_loc=first.dst_loc, + src_indices=torch.cat([unit.src_indices for unit in units]), + dst_indices=torch.cat([unit.dst_indices for unit in units]), + op_id=-1, + is_retract=any(unit.is_retract for unit in units), + layerwise_cow_src_indices=( + torch.cat(cow_src_indices) if cow_src_indices else None + ), + layerwise_cow_dst_indices=( + torch.cat(cow_dst_indices) if cow_dst_indices else None + ), + ) + + def _prepare_indices( + self, unit: TransferUnit, pool: CachePool + ) -> tuple[torch.Tensor, torch.Tensor]: + if unit.src_loc == Location.HOST: + host_indices = unit.src_indices + device_indices = unit.dst_indices + elif unit.dst_loc == Location.HOST: + host_indices = unit.dst_indices + device_indices = unit.src_indices + else: + raise ValueError(f"unsupported transfer direction: {unit.direction}") + + io_backend = getattr(pool, "io_backend", self.io_backend) + if io_backend == "kernel": + target_device = pool.device + if device_indices.device != target_device: + device_indices = device_indices.to(target_device, non_blocking=True) + if host_indices.device != target_device: + host_indices = host_indices.to(target_device, non_blocking=True) + elif io_backend == "direct": + if pool.host_layout == "layer_first": + device_indices = device_indices.cpu() + host_indices, idx = host_indices.sort() + device_indices = device_indices.index_select(0, idx) + else: + raise ValueError(f"Unsupported host layout: {pool.host_layout}") + else: + raise ValueError(f"Unsupported io_backend={io_backend}") + + if unit.src_loc == Location.HOST: + return host_indices, device_indices + return device_indices, host_indices + + def _pool_writeback( + self, pool: CachePool, src_indices: torch.Tensor, dst_indices: torch.Tensor + ) -> None: + try: + pool.writeback( + src_indices, dst_indices, block_quota=self._writeback_block_quota + ) + except TypeError as exc: + if "block_quota" not in str(exc): + raise + pool.writeback(src_indices, dst_indices) + + @staticmethod + def _record_if_cuda(tensor: torch.Tensor, stream) -> None: + if tensor.is_cuda: + tensor.record_stream(stream) + + def drain(self) -> list: + results: list = [] + results.extend(self._poll_write_acks()) + results.extend(self._poll_load_acks()) + return results + + def _poll_write_acks(self) -> list: + results = [] + completed_writebacks = getattr(self, "completed_writebacks", []) + for op_id in completed_writebacks: + logger.debug("[cache_op] writeback done op_id=%s immediate=True", op_id) + evt = Cache.WriteBackDoneEvent() + evt.op_id = op_id + evt.success = True + results.append(evt) + completed_writebacks.clear() + remaining = [] + for ack in self.ack_write_queue: + if ack.finish_event.query(): + logger.debug( + "[cache_op] writeback done op_ids=%s immediate=False", ack.op_ids + ) + for op_id in ack.op_ids: + evt = Cache.WriteBackDoneEvent() + evt.op_id = op_id + evt.success = True + results.append(evt) + else: + remaining.append(ack) + self.ack_write_queue[:] = remaining + return results + + def _poll_load_acks(self) -> list: + results = [] + remaining = [] + for ack in self.ack_load_queue: + if not ack.finish_event.query(): + remaining.append(ack) + self.ack_load_queue[:] = remaining + return results + + def get_producer_index( + self, kind_or_op_id: CacheKind | str | int, op_id: int | None = None + ) -> int | None: + if op_id is None: + kind = CacheKind.KV + op_id = int(kind_or_op_id) + else: + kind = CacheKind(kind_or_op_id) + return self._producer_map[kind].pop(int(op_id), None) + + def set_consumer( + self, + kind_or_producer_index: CacheKind | str | int | Iterable[int], + producer_index: int | Iterable[int] | None = None, + ) -> None: + if producer_index is None: + kind = CacheKind.KV + producer_index = kind_or_producer_index + else: + kind = CacheKind(kind_or_producer_index) + self._counters[kind].set_consumer(producer_index) + + def shutdown(self) -> None: + self.write_stream.synchronize() + self.load_stream.synchronize() + for pool in self.pools.values(): + shutdown = getattr(pool, "shutdown", None) + if shutdown is not None: + shutdown() + + def reset(self) -> None: + self.write_stream.synchronize() + self.load_stream.synchronize() + self._clear_queues(self.write_queues) + self._clear_queues(self.load_queues) + self.ack_write_queue.clear() + self.ack_load_queue.clear() + for producer_map in self._producer_map.values(): + producer_map.clear() + for counter in self._counters.values(): + counter.reset() diff --git a/python/tokenspeed/runtime/cache/executor/memory_executor.py b/python/tokenspeed/runtime/cache/executor/memory_executor.py new file mode 100644 index 0000000..98c037b --- /dev/null +++ b/python/tokenspeed/runtime/cache/executor/memory_executor.py @@ -0,0 +1,500 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Top-level memory executor that coordinates host and storage executors.""" + +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import dataclass + +from tokenspeed_scheduler import Cache + +from tokenspeed.runtime.cache.executor.host_executor import HostExecutor +from tokenspeed.runtime.cache.executor.storage_executor import StorageExecutor +from tokenspeed.runtime.cache.kv_cache_host import ( + DSATokenToKVPoolHost, + MHATokenToKVPoolHost, + MLATokenToKVPoolHost, + get_available_host_memory_bytes, +) +from tokenspeed.runtime.cache.mamba_cache_host import MambaPoolHost +from tokenspeed.runtime.cache.transfer.kv_pool import KVCachePool +from tokenspeed.runtime.cache.transfer.mamba_pool import MambaCachePool +from tokenspeed.runtime.cache.transfer.types import CacheKind +from tokenspeed.runtime.layers.attention.kv_cache.dsa import DSATokenToKVPool +from tokenspeed.runtime.layers.attention.kv_cache.mha import MHATokenToKVPool +from tokenspeed.runtime.layers.attention.kv_cache.mla import MLATokenToKVPool +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +@dataclass(slots=True) +class MemoryExecutorConfig: + layer_num: int + page_size: int = 64 + host_ratio: float = 2.0 + host_size_gb: int = 0 + host_parallel_count: int = 1 + host_reserve_gb: float = 10.0 + io_backend: str = "kernel" + host_layout: str = "layer_first" + storage_backend: str | None = "mooncake" + storage_backend_extra_config: str | None = None + model_name: str | None = None + enable_mamba_l2: bool = False + mamba_l2_host_slots: int = 0 + mamba_l2_layout: str = "layer_first" + mamba_l2_io_backend: str = "kernel" + + +def _aligned_token_count(size: int, page_size: int) -> int: + return (size // page_size + 1) * page_size + + +def _pool_size_per_token(pool) -> int: + dtype_size = pool.store_dtype.itemsize + if isinstance(pool, DSATokenToKVPool): + latent_size = ( + (pool.kv_lora_rank + pool.qk_rope_head_dim) * dtype_size * pool.layer_num + ) + return latent_size + pool.index_k_row_bytes * pool.layer_num + if isinstance(pool, MHATokenToKVPool): + return pool.head_dim * pool.head_num * pool.layer_num * dtype_size * 2 + if isinstance(pool, MLATokenToKVPool): + return (pool.kv_lora_rank + pool.qk_rope_head_dim) * dtype_size * pool.layer_num + raise ValueError(f"Unsupported KV pool type for host budget: {type(pool)}") + + +def _auto_capped_host_size_tokens( + *, + requested_tokens: int, + page_size: int, + size_per_token: int, + available_host_memory_bytes: int, + host_parallel_count: int, +) -> int: + """Return a HostKVCache host_size_tokens override, or 0 when no cap is needed.""" + + aligned_requested_tokens = _aligned_token_count(requested_tokens, page_size) + requested_bytes = aligned_requested_tokens * size_per_token + per_rank_budget = available_host_memory_bytes // max(host_parallel_count, 1) + if requested_bytes <= per_rank_budget: + return 0 + + budget_tokens = per_rank_budget // max(size_per_token, 1) + max_aligned_tokens = (budget_tokens // page_size) * page_size + if max_aligned_tokens <= page_size: + raise ValueError( + "Not enough host memory available for KVStore after cgroup-aware " + f"budgeting: per-rank budget={per_rank_budget / 1e9:.2f} GB, " + f"size_per_token={size_per_token}." + ) + return max_aligned_tokens - page_size + + +class MemoryExecutor: + """Coordinate host-memory and storage-backed cache operations.""" + + def __init__( + self, + device_pool, + config: MemoryExecutorConfig, + is_dp_attention_enabled: bool, + tp_group=None, + draft_device_pool=None, + mamba_pool=None, + ): + self.page_size = config.page_size + kv_pool_types = (DSATokenToKVPool, MHATokenToKVPool, MLATokenToKVPool) + + # Unwrap LayerMappedKVPool (hybrid GDN models) to get the inner MHA pool. + actual_pool = device_pool + if hasattr(device_pool, "inner") and not isinstance(device_pool, kv_pool_types): + actual_pool = device_pool.inner + + actual_draft_pool = None + if draft_device_pool is not None: + actual_draft_pool = draft_device_pool + if hasattr(draft_device_pool, "inner") and not isinstance( + draft_device_pool, kv_pool_types + ): + actual_draft_pool = draft_device_pool.inner + if not isinstance(actual_draft_pool, kv_pool_types): + raise ValueError( + f"draft_device_pool only supports DSA, MHA and MLA, " + f"got {type(actual_draft_pool)}" + ) + + host_size_tokens = 0 + if config.host_size_gb == 0: + target_size_per_token = _pool_size_per_token(actual_pool) + draft_size_per_token = ( + _pool_size_per_token(actual_draft_pool) + if actual_draft_pool is not None + else 0 + ) + combined_size_per_token = target_size_per_token + draft_size_per_token + reserve_bytes = int(config.host_reserve_gb * (1024**3)) + available_bytes, _, cgroup_available = get_available_host_memory_bytes( + reserve_bytes + ) + requested_tokens = int(actual_pool.size * config.host_ratio) + host_size_tokens = _auto_capped_host_size_tokens( + requested_tokens=requested_tokens, + page_size=config.page_size, + size_per_token=combined_size_per_token, + available_host_memory_bytes=available_bytes, + host_parallel_count=config.host_parallel_count, + ) + if host_size_tokens > 0: + capped_tokens = _aligned_token_count(host_size_tokens, config.page_size) + requested_tokens_aligned = _aligned_token_count( + requested_tokens, config.page_size + ) + logger.warning( + "Capping KVStore host pool for cgroup budget: " + "tokens %s -> %s, total bytes %.2f GB -> %.2f GB " + "(parallel_count=%s, available=%.2f GB, cgroup_available=%s)", + requested_tokens_aligned, + capped_tokens, + requested_tokens_aligned * combined_size_per_token / 1e9, + capped_tokens * combined_size_per_token / 1e9, + config.host_parallel_count, + available_bytes / 1e9, + ( + f"{cgroup_available / 1e9:.2f} GB" + if cgroup_available is not None + else "unlimited" + ), + ) + + # DSA subclasses MLA, so it must be matched before the MLA branch. + if isinstance(actual_pool, DSATokenToKVPool): + self.host_pool = DSATokenToKVPoolHost( + actual_pool, + config.host_ratio, + config.host_size_gb, + config.page_size, + config.host_layout, + host_size_tokens=host_size_tokens, + ) + elif isinstance(actual_pool, MHATokenToKVPool): + self.host_pool = MHATokenToKVPoolHost( + actual_pool, + config.host_ratio, + config.host_size_gb, + config.page_size, + config.host_layout, + host_size_tokens=host_size_tokens, + ) + elif isinstance(actual_pool, MLATokenToKVPool): + self.host_pool = MLATokenToKVPoolHost( + actual_pool, + config.host_ratio, + config.host_size_gb, + config.page_size, + config.host_layout, + host_size_tokens=host_size_tokens, + ) + else: + raise ValueError( + f"host_pool only supports DSA, MHA and MLA, got {type(actual_pool)} " + f"from module {type(actual_pool).__module__}" + ) + + # Draft model L2 cache: draft shares the same page mapping as the base + # model, so its host pool must hold exactly the same number of tokens. + # Pass host_size_tokens directly to bypass ratio/GB recalculation. + if actual_draft_pool is not None: + if isinstance(actual_draft_pool, DSATokenToKVPool): + self.draft_host_pool = DSATokenToKVPoolHost( + actual_draft_pool, + config.host_ratio, + config.host_size_gb, + config.page_size, + config.host_layout, + host_size_tokens=self.host_pool.size, + ) + elif isinstance(actual_draft_pool, MHATokenToKVPool): + self.draft_host_pool = MHATokenToKVPoolHost( + actual_draft_pool, + config.host_ratio, + config.host_size_gb, + config.page_size, + config.host_layout, + host_size_tokens=self.host_pool.size, + ) + elif isinstance(actual_draft_pool, MLATokenToKVPool): + self.draft_host_pool = MLATokenToKVPoolHost( + actual_draft_pool, + config.host_ratio, + config.host_size_gb, + config.page_size, + config.host_layout, + host_size_tokens=self.host_pool.size, + ) + else: + raise ValueError( + f"draft_device_pool only supports DSA, MHA and MLA, " + f"got {type(actual_draft_pool)}" + ) + draft_host_bytes = ( + self.draft_host_pool.size * self.draft_host_pool.size_per_token + ) + logger.info( + "Allocating %.2f GB host memory for draft model L2 cache (pool_type=%s size_tokens=%s size_per_token=%s layer_num=%s)", + draft_host_bytes / 1e9, + type(self.draft_host_pool).__name__, + self.draft_host_pool.size, + self.draft_host_pool.size_per_token, + actual_draft_pool.layer_num, + ) + draft_layer_num = actual_draft_pool.layer_num + else: + self.draft_host_pool = None + draft_layer_num = 0 + + pools = None + self.mamba_host_pool = None + if ( + config.enable_mamba_l2 + and mamba_pool is not None + and config.mamba_l2_host_slots > 0 + ): + self.mamba_host_pool = MambaPoolHost( + mamba_pool, + host_size_slots=config.mamba_l2_host_slots, + layout=config.mamba_l2_layout, + ) + pools = [ + KVCachePool( + device_pool=device_pool, + host_pool=self.host_pool, + io_backend=config.io_backend, + layer_num=actual_pool.layer_num, + draft_device_pool=( + actual_draft_pool if draft_device_pool is not None else None + ), + draft_host_pool=self.draft_host_pool, + draft_layer_num=draft_layer_num, + ), + MambaCachePool( + device_pool=mamba_pool, + host_pool=self.mamba_host_pool, + io_backend=config.mamba_l2_io_backend, + ), + ] + logger.debug( + "[cache_op] MemoryExecutor init pools=%s host_pools=%s draft=%s mamba=%s io_backend=%s host_layout=%s", + [pool.kind.value for pool in pools], + [type(self.host_pool).__name__, type(self.mamba_host_pool).__name__], + self.draft_host_pool is not None, + True, + config.io_backend, + config.host_layout, + ) + + if pools is not None: + self.host_exec = HostExecutor(pools=pools, io_backend=config.io_backend) + else: + self.host_exec = HostExecutor( + page_size=config.page_size, + device_pool=device_pool, + host_pool=self.host_pool, + io_backend=config.io_backend, + layer_num=actual_pool.layer_num, + draft_device_pool=( + actual_draft_pool if draft_device_pool is not None else None + ), + draft_host_pool=self.draft_host_pool, + draft_layer_num=draft_layer_num, + ) + self.storage_exec = StorageExecutor( + page_size=config.page_size, + device_pool=device_pool, + host_pool=self.host_pool, + storage_backend_type=config.storage_backend, + storage_backend_extra_config=config.storage_backend_extra_config, + model_name=config.model_name, + is_dp_attention_enabled=is_dp_attention_enabled, + tp_group=tp_group, + ) + self._pending_mamba_layerwise_cow: dict[int, list[int]] | None = None + + @staticmethod + def _page_groups_by_kind(op) -> dict[CacheKind, tuple[list, list]]: + src_by_kind = getattr(op, "src_pages_by_kind", None) + dst_by_kind = getattr(op, "dst_pages_by_kind", None) + if src_by_kind is None or dst_by_kind is None: + return {CacheKind.KV: (op.src_pages, op.dst_pages)} + groups: dict[CacheKind, tuple[list, list]] = {} + for kind in CacheKind: + src_pages = src_by_kind.get(kind.value, []) + dst_pages = dst_by_kind.get(kind.value, []) + groups[kind] = (src_pages, dst_pages) + return groups + + def set_mamba_layerwise_cow( + self, cow_dst_pages_by_src: dict[int, list[int]] | None + ) -> None: + self._pending_mamba_layerwise_cow = cow_dst_pages_by_src or None + + def submit_plan(self, plan) -> None: + if plan.cache: + logger.debug("[cache_op] submit_plan: %s cache ops", len(plan.cache)) + try: + for op in plan.cache: + self.submit(op) + self.host_exec.flush() + finally: + self._pending_mamba_layerwise_cow = None + + def submit(self, op) -> None: + if isinstance(op, Cache.WriteBackOp): + logger.debug( + "[cache_op] writeback op_id=%s src_pages=%s dst_pages=%s", + op.op_ids, + len(op.src_pages), + len(op.dst_pages), + ) + groups = self._page_groups_by_kind(op) + for i in range(len(op.op_ids)): + op_id = op.op_ids[i] + is_retract = bool(getattr(op, "is_retract", [False])[i]) + for kind, (src_groups, dst_groups) in groups.items(): + if kind not in self.host_exec.pools: + continue + src_pages = src_groups[i] if i < len(src_groups) else [] + dst_pages = dst_groups[i] if i < len(dst_groups) else [] + if not src_pages: + continue + if kind == CacheKind.MAMBA: + logger.debug( + "[cache_op][mamba_l2] writeback schedule " + "op_id=%s slots=%s device_slots=%s host_slots=%s " + "is_retract=%s", + op_id, + len(src_pages), + src_pages[:8], + dst_pages[:8], + is_retract, + ) + self.host_exec.enqueue_writeback( + op_id, + src_pages, + dst_pages, + is_retract=is_retract, + kind=kind, + ) + if all( + i >= len(src_groups) or not src_groups[i] + for kind, (src_groups, _) in groups.items() + if kind in self.host_exec.pools + ): + self.host_exec.completed_writebacks.append(op_id) + elif isinstance(op, Cache.LoadBackOp): + logger.debug( + "[cache_op] loadback op_id=%s src_pages=%s dst_pages=%s", + op.op_ids, + len(op.src_pages), + len(op.dst_pages), + ) + groups = self._page_groups_by_kind(op) + for i in range(len(op.op_ids)): + op_id = op.op_ids[i] + for kind, (src_groups, dst_groups) in groups.items(): + if kind not in self.host_exec.pools: + continue + src_pages = src_groups[i] if i < len(src_groups) else [] + dst_pages = dst_groups[i] if i < len(dst_groups) else [] + if not src_pages: + continue + if kind == CacheKind.MAMBA: + logger.debug( + "[cache_op][mamba_l2] loadback schedule " + "op_id=%s slots=%s host_slots=%s device_slots=%s", + op_id, + len(src_pages), + src_pages[:8], + dst_pages[:8], + ) + loadback_kwargs = {} + mamba_layerwise_cow = getattr( + self, "_pending_mamba_layerwise_cow", None + ) + if kind == CacheKind.MAMBA and mamba_layerwise_cow: + loadback_kwargs["layerwise_cow_dst_pages_by_src"] = ( + mamba_layerwise_cow + ) + self.host_exec.enqueue_loadback( + op_id, src_pages, dst_pages, kind=kind, **loadback_kwargs + ) + + elif isinstance(op, Cache.PrefetchOp): + logger.debug( + "[cache_op] prefetch op_id=%s dst_pages=%s", op.op_id, len(op.dst_pages) + ) + self.storage_exec.submit_prefetch(op) + elif isinstance(op, Cache.BackUpOp): + logger.debug( + "[cache_op] backup op_id=%s src_pages=%s", op.op_id, len(op.src_pages) + ) + self.storage_exec.submit_backup(op) + else: + raise ValueError("unsupported cache op kind") + + def poll_results(self) -> list: + results: list = [] + results.extend(self.host_exec.drain()) + results.extend(self.storage_exec.drain()) + if results: + for r in results: + logger.debug( + "[cache_op] done op_id=%s success=%s type=%s", + r.op_id, + r.success, + type(r).__name__, + ) + return results + + def get_producer_index( + self, kind_or_op_id: CacheKind | str | int, op_id: int | None = None + ) -> int | None: + return self.host_exec.get_producer_index(kind_or_op_id, op_id) + + def set_consumer( + self, + kind_or_producer_index: CacheKind | str | int | Iterable[int], + producer_index: int | Iterable[int] | None = None, + ) -> None: + self.host_exec.set_consumer(kind_or_producer_index, producer_index) + + def query_l3_pages(self, hashes: list[str]) -> int: + return self.storage_exec.query_exists(hashes) + + def shutdown(self) -> None: + self.host_exec.shutdown() + self.storage_exec.shutdown() + + def reset(self) -> None: + self.host_exec.reset() + self.storage_exec.drain() diff --git a/python/tokenspeed/runtime/cache/executor/storage_executor.py b/python/tokenspeed/runtime/cache/executor/storage_executor.py new file mode 100644 index 0000000..35b6661 --- /dev/null +++ b/python/tokenspeed/runtime/cache/executor/storage_executor.py @@ -0,0 +1,487 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Storage-backed executor for cache prefetch and backup operations.""" + +from __future__ import annotations + +import json +import threading +import time +from concurrent.futures import Future, ThreadPoolExecutor +from functools import partial +from queue import Empty, Queue + +import torch +import torch.distributed as dist +from tokenspeed_scheduler import Cache + +from tokenspeed.runtime.cache.executor.host_executor import ( + page_ids_to_token_indices, +) +from tokenspeed.runtime.cache.kvstore_storage import KVStoreStorageConfig +from tokenspeed.runtime.cache.storage import StorageBackendFactory +from tokenspeed.runtime.distributed.process_group_manager import _make_all_groups +from tokenspeed.runtime.layers.attention.kv_cache.mla import MLATokenToKVPool +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +def _parse_storage_backend_extra_config(raw: str | None): + extra_config = {} + if raw: + try: + extra_config = json.loads(raw) + except json.JSONDecodeError as exc: + logger.error("Invalid backend extra config JSON: %s", exc) + raise + + prefetch_threshold = extra_config.pop("prefetch_threshold", 256) + prefetch_timeout_base = extra_config.pop("prefetch_timeout_base", 1) + prefetch_timeout_per_ki_token = extra_config.pop( + "prefetch_timeout_per_ki_token", 0.25 + ) + + if not isinstance(prefetch_threshold, int): + raise ValueError( + f"prefetch_threshold must be int, got {type(prefetch_threshold).__name__}" + ) + if not isinstance(prefetch_timeout_base, (int, float)): + raise ValueError( + f"prefetch_timeout_base must be number, got {type(prefetch_timeout_base).__name__}" + ) + if not isinstance(prefetch_timeout_per_ki_token, (int, float)): + raise ValueError( + f"prefetch_timeout_per_ki_token must be number, got {type(prefetch_timeout_per_ki_token).__name__}" + ) + + return ( + extra_config, + prefetch_threshold, + float(prefetch_timeout_base), + float(prefetch_timeout_per_ki_token), + ) + + +def _generate_storage_config( + device_pool, + host_pool, + model_name: str | None, + extra_config_dict: dict, + is_dp_attention_enabled: bool, +) -> KVStoreStorageConfig: + tp_rank = dist.get_rank() if dist.is_initialized() else 0 + tp_size = dist.get_world_size() if dist.is_initialized() else 1 + + is_mla_backend = isinstance(device_pool, MLATokenToKVPool) + + return KVStoreStorageConfig( + tp_rank=tp_rank, + tp_size=tp_size, + is_mla_model=is_mla_backend, + is_page_first_layout=host_pool.layout == "page_first", + model_name=model_name, + extra_config=extra_config_dict, + ) + + +class StorageExecutor: + """Execute L3 storage prefetch and backup operations asynchronously.""" + + def __init__( + self, + page_size: int, + device_pool, + host_pool, + storage_backend_type: str | None, + storage_backend_extra_config: str | None = None, + model_name: str | None = None, + is_dp_attention_enabled: bool = False, + storage_batch_size: int = 128, + tp_group=None, + ): + self.page_size = page_size + self.host_pool = host_pool + self.storage_batch_size = storage_batch_size + ( + extra_config_dict, + _prefetch_threshold, + prefetch_timeout_base, + prefetch_timeout_per_ki_token, + ) = _parse_storage_backend_extra_config(storage_backend_extra_config) + + self.prefetch_timeout_base = prefetch_timeout_base + self.prefetch_timeout_per_page = ( + page_size / 1024 * prefetch_timeout_per_ki_token + ) + + self.storage_backend = None + if storage_backend_type is not None: + storage_config = _generate_storage_config( + device_pool, + host_pool, + model_name, + extra_config_dict, + is_dp_attention_enabled, + ) + try: + self.storage_backend = StorageBackendFactory.create_backend( + storage_backend_type, storage_config, host_pool + ) + except ValueError as exc: + raise ValueError(f"Failed to create storage backend: {exc}") from exc + + self.storage_backend.register_mem_pool_host(host_pool) + self.tp_size = ( + torch.distributed.get_world_size(group=tp_group) + if tp_group is not None + else 1 + ) + # Dedicated subgroup for kvstore collectives. The shared ``tp_group`` + # is used by the engine main thread elsewhere (event_loop, + # request_handler). Issuing collectives on it from + # the aggregator thread would race those callers — different threads + # on the same rank issuing on the same group can pair up out-of-order + # across ranks and hang. + # + # ``new_group`` is itself a world-wide collective, so when there is + # more than one TP-shaped group in the world (e.g. DP attention with + # attn_tp groups [0..3] and [4..7]) we have to call it for *every* + # such group in the same deterministic order on every rank — calling + # it only with the local rank's TP group would deadlock other ranks. + # ``_make_all_groups`` enumerates the full set with the same size and + # stride pattern, mirroring ``pg_manager.init_process_group``. + self._tp_group = None + if tp_group is not None and self.tp_size > 1: + my_ranks = tuple(torch.distributed.get_process_group_ranks(tp_group)) + for g in _make_all_groups(my_ranks): + pg = torch.distributed.new_group(ranks=list(g), backend="gloo") + if g == my_ranks: + self._tp_group = pg + self._results: Queue = Queue() + self._prefetch_op_to_rid: dict = {} # op_id → request_id + self._executor: ThreadPoolExecutor | None = None + # All collectives on the dedicated kvstore subgroup are funneled + # through a single aggregator thread so they issue in deterministic + # submit order across ranks (Gloo and NCCL groups both require + # callers to agree on issuance order; concurrent issuance from + # worker threads can deadlock when ops finish in different order on + # each rank). + self._aggregator_pending: Queue = Queue() + self._aggregator_stop = threading.Event() + self._aggregator_thread: threading.Thread | None = None + if self.storage_backend is not None: + self._executor = ThreadPoolExecutor( + max_workers=2, + thread_name_prefix="tokenspeed-mem-l3-io", + ) + self._aggregator_thread = threading.Thread( + target=self._aggregator_loop, + name="tokenspeed-mem-l3-aggr", + daemon=True, + ) + self._aggregator_thread.start() + + @property + def enabled(self) -> bool: + return self.storage_backend is not None + + def submit_prefetch(self, op) -> None: + # Extract request_id from the op and remember mapping + rid = op.request_id + self._prefetch_op_to_rid[op.op_id] = rid + + if not self.enabled: + evt = Cache.PrefetchDoneEvent() + evt.op_id = op.op_id + evt.request_id = rid + evt.success = False + evt.completed_pages = 0 + self._results.put(evt) + return + future = self._executor.submit(self._run_prefetch, op) + # Enqueue at submit time (not completion time) so both ranks see the + # same aggregator order, which guarantees the per-op TP all_reduce + # pairs up correctly. + self._aggregator_pending.put(("prefetch", op.op_id, rid, future)) + + def submit_backup(self, op) -> None: + if not self.enabled: + evt = Cache.BackUpDoneEvent() + evt.op_id = op.op_id + evt.success = False + self._results.put(evt) + return + future = self._executor.submit(self._run_backup, op) + future.add_done_callback(partial(self._on_backup_done, op.op_id)) + + def _prefetch_deadline(self, n_pages: int) -> float: + return ( + time.monotonic() + + self.prefetch_timeout_base + + n_pages * self.prefetch_timeout_per_page + ) + + def _run_prefetch(self, op) -> int: + hashes = op.rolling_page_hashes + if not hashes: + logger.debug("[cache_op] prefetch_exec op_id=%s no hashes, skip", op.op_id) + return 0 + dst_pages = op.dst_pages + if len(hashes) != len(dst_pages): + raise ValueError( + f"prefetch key/page mismatch: {len(hashes)} hashes " + f"vs {len(dst_pages)} dst_pages" + ) + + deadline = self._prefetch_deadline(len(hashes)) + completed_pages = 0 + + for i in range(0, len(hashes), self.storage_batch_size): + if time.monotonic() > deadline: + logger.warning( + "prefetch op %s timed out after %s/%s pages", + op.op_id, + completed_pages, + len(hashes), + ) + break + + batch_hashes = hashes[i : i + self.storage_batch_size] + batch_dst = dst_pages[i : i + len(batch_hashes)] + host_indices = page_ids_to_token_indices(batch_dst, self.page_size, "cpu") + try: + results = self.storage_backend.batch_get_v1(batch_hashes, host_indices) + except Exception as exc: + logger.warning( + "prefetch op %s batch IO error at offset %s: %s", op.op_id, i, exc + ) + break + result_ok = 0 + for ok in results: + if not ok: + break + result_ok += 1 + completed_pages += result_ok + if result_ok < len(results): + logger.warning( + "prefetch op %s: %s/%s pages missed in batch at offset %s", + op.op_id, + len(results) - result_ok, + len(results), + i, + ) + break + + # NOTE: TP all_reduce moved to the aggregator thread; the worker only + # performs storage I/O and returns the local count. See _aggregator_loop. + logger.debug( + "[cache_op] prefetch_exec op_id=%s completed %s/%s pages (local)", + op.op_id, + completed_pages, + len(hashes), + ) + return completed_pages + + def _run_backup(self, op) -> None: + hashes = op.rolling_page_hashes + src_pages = op.src_pages + total_num_pages = len(hashes) + logger.debug( + "[cache_op] backup_exec op_id=%s pages=%s", op.op_id, total_num_pages + ) + completed_pages = 0 + for i in range(0, total_num_pages, self.storage_batch_size): + batch_hashes = hashes[i : i + self.storage_batch_size] + batch_src = src_pages[i : i + len(batch_hashes)] + host_indices = page_ids_to_token_indices(batch_src, self.page_size, "cpu") + results = self.storage_backend.batch_set_v1(batch_hashes, host_indices) + result_ok = sum(1 for ok in results if ok) + completed_pages += result_ok + if result_ok < len(results): + failed_count = len(results) - result_ok + logger.warning( + "backup op %s: %s/%s pages failed in batch at offset %s", + op.op_id, + failed_count, + len(results), + i, + ) + raise RuntimeError( + f"backup op {op.op_id}: {completed_pages}/{total_num_pages} pages ok, " + f"{failed_count} failed in batch at offset {i}" + ) + logger.debug( + "[cache_op] backup_exec op_id=%s done, all %s/%s pages ok", + op.op_id, + completed_pages, + total_num_pages, + ) + + def _aggregator_loop(self) -> None: + """Single thread that owns all collectives on ``self._tp_group``. + + ``self._tp_group`` is a dedicated kvstore-only Gloo subgroup, separate + from the ``tp_group`` used by the main thread elsewhere in the engine + — so this thread's collectives don't race main-thread collectives. + + Both prefetch completions (which need a TP MIN all_reduce on + ``completed_pages``) and ``query_exists`` calls are funneled through + ``_aggregator_pending`` in submit order. Because both ranks observe + the same submit order (deterministic from the C++ scheduler + + main-thread engine loop), the aggregator on each rank issues + collectives in the same order, so the MIN all_reduces pair correctly + across ranks. + """ + while not self._aggregator_stop.is_set(): + try: + item = self._aggregator_pending.get(block=True, timeout=1) + except Empty: + continue + kind = item[0] + if kind == "prefetch": + _, op_id, rid, future = item + evt = Cache.PrefetchDoneEvent() + evt.op_id = op_id + evt.request_id = self._prefetch_op_to_rid.pop(op_id, rid) + # Encode local outcome with a -1 sentinel for "this rank + # failed locally". A bare ``continue`` on local failure would + # skip the all_reduce below and deadlock peers that *did* + # succeed locally — the collective must run in lockstep on + # every rank. ReduceOp.MIN propagates the sentinel across all + # ranks, so every rank agrees the op failed. + try: + local = future.result() + except Exception as exc: + logger.error("prefetch op %s local I/O failed: %s", op_id, exc) + local = -1 + completed_pages = local + if self.tp_size > 1: + try: + t = torch.tensor(local, dtype=torch.int) + torch.distributed.all_reduce( + t, + op=torch.distributed.ReduceOp.MIN, + group=self._tp_group, + ) + completed_pages = t.item() + except Exception as exc: + logger.warning("prefetch op %s TP sync failed: %s", op_id, exc) + completed_pages = -1 + if completed_pages < 0: + evt.success = False + evt.completed_pages = 0 + else: + evt.success = True + evt.completed_pages = completed_pages + logger.debug( + "[prefetch_done] op_id=%s request_id=%s success=%s completed_pages=%s", + op_id, + evt.request_id, + evt.success, + evt.completed_pages, + ) + self._results.put(evt) + elif kind == "query_exists": + _, hashes, result_future = item + # Same reasoning as the prefetch branch: never short-circuit + # before the collective. + try: + local = self._query_exists_local(hashes) + except Exception as exc: + logger.error("query_exists local I/O failed: %s", exc) + local = -1 + total_hit = local + if self.tp_size > 1: + try: + t = torch.tensor(local, dtype=torch.int) + torch.distributed.all_reduce( + t, + op=torch.distributed.ReduceOp.MIN, + group=self._tp_group, + ) + total_hit = t.item() + except Exception as exc: + logger.warning("query_exists TP sync failed: %s", exc) + total_hit = -1 + if total_hit < 0: + result_future.set_exception( + RuntimeError("query_exists failed on at least one rank") + ) + else: + result_future.set_result(total_hit) + else: + logger.error("unknown aggregator item kind: %s", kind) + + def _on_backup_done(self, op_id: int, future) -> None: + evt = Cache.BackUpDoneEvent() + evt.op_id = op_id + try: + future.result() + evt.success = True + except Exception as exc: + evt.success = False + logger.error("backup op %s failed: %s", op_id, exc) + self._results.put(evt) + + def drain(self) -> list: + results = [] + while True: + try: + results.append(self._results.get_nowait()) + except Empty: + return results + + def query_exists(self, hashes: list[str]) -> int: + if not self.enabled or not hashes: + return 0 + if self.tp_size <= 1 or self._aggregator_thread is None: + return self._query_exists_local(hashes) + # Route through the aggregator so the all_reduce on ``tp_group`` is + # serialized with prefetch-completion all_reduces on the same group. + result_future: Future = Future() + self._aggregator_pending.put(("query_exists", list(hashes), result_future)) + return result_future.result() + + def _query_exists_local(self, hashes: list[str]) -> int: + total_hit = 0 + for i in range(0, len(hashes), self.storage_batch_size): + batch = hashes[i : i + self.storage_batch_size] + hit = self.storage_backend.batch_exists(batch) + total_hit += hit + if hit < len(batch): + break + return total_hit + + def shutdown(self) -> None: + if self._aggregator_thread is not None: + self._aggregator_stop.set() + self._aggregator_thread.join(timeout=10) + self._aggregator_thread = None + if self._executor is not None: + self._executor.shutdown(wait=True) + self._executor = None + if self.storage_backend is not None and hasattr(self.storage_backend, "close"): + try: + self.storage_backend.close() + except Exception: + logger.exception("Failed to close storage backend") + self.storage_backend = None diff --git a/python/tokenspeed/runtime/cache/flat_host_mirror.py b/python/tokenspeed/runtime/cache/flat_host_mirror.py new file mode 100644 index 0000000..504f3d4 --- /dev/null +++ b/python/tokenspeed/runtime/cache/flat_host_mirror.py @@ -0,0 +1,218 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Byte-blind pinned-CPU mirror of a device KV pool for the flat L2 host +tier (M15 Phase D). Transport mechanism only; scheduler/engine wiring is D2. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Sequence + +import torch + + +def _identity_dedup( + tensors: Sequence[torch.Tensor | None], +) -> list[torch.Tensor]: + """Distinct tensors in first-appearance order; None slots (flat GDN + state layers carry no KV) are skipped.""" + seen: dict[int, torch.Tensor] = {} + for t in tensors: + if t is None: + continue + seen.setdefault(id(t), t) + return list(seen.values()) + + +def _state_slabs(device_kv_pool) -> list[tuple[torch.Tensor, torch.Tensor]]: + """(conv, ssm) state slab pairs, [] on pools predating state slabs.""" + return list(getattr(device_kv_pool, "state_slabs", None) or ()) + + +def flat_bytes_per_host_page(device_kv_pool) -> int: + """Bytes one host page occupies across all mirrors, computed from the + device pool alone (no mirror allocation) -- the sizing side of + ``FlatHostMirror.bytes_per_host_page`` for host-budget arithmetic. + """ + tensors = _identity_dedup(device_kv_pool.k_buffer) + _identity_dedup( + device_kv_pool.v_buffer + ) + page_size = int(device_kv_pool.page_size) + kv_bytes = sum(t.element_size() * t[0].numel() * page_size for t in tensors) + # State slabs are page-indexed: one constant row per page id. + state_bytes = sum( + t.element_size() * t[0].numel() + for pair in _state_slabs(device_kv_pool) + for t in pair + ) + return kv_bytes + state_bytes + + +class FlatHostMirror: + """One pinned CPU mirror per DISTINCT device KV tensor plus one per + state slab tensor; a (device_page, host_page) pair copies that page's + row range on every mirror pair. + + Slab tensors are enumerated once each -- a page's rows are exactly its + owner group's layers, so byte copies are group-safe by id-exclusivity. + + ``tensor_pairs`` order (PINNED, D2 fencing indexes into it): K*, V*, + then state tensors flattened in slab order (conv0, ssm0, conv1, ...). + KV mirrors span ``page_size`` token rows per page; state slabs are + page-indexed (one snapshot row per page id), so their mirrors span 1 + row per page -- ``row_spans[i]`` carries each pair's span. + """ + + def __init__(self, device_kv_pool, num_host_pages: int): + self.page_size = int(device_kv_pool.page_size) + self.num_host_pages = int(num_host_pages) + + # Slab layout dedups the per-layer entries to one K + one V slab per + # paired layer set (layers-per-group slabs); legacy layout keeps all + # per-layer buffers (dead-row copies are harmless). + k_tensors = _identity_dedup(device_kv_pool.k_buffer) + v_tensors = _identity_dedup(device_kv_pool.v_buffer) + self.num_k_tensors = len(k_tensors) + + k_index = {id(t): i for i, t in enumerate(k_tensors)} + v_index = {id(t): i for i, t in enumerate(v_tensors)} + # None entries (flat GDN state layers, no KV) map to None: those + # layers fence on state_tensor_indices_of_layer instead. + self._layer_to_k_index = [ + None if t is None else k_index[id(t)] for t in device_kv_pool.k_buffer + ] + # Invariant D2 relies on: a layer's V tensor sits at + # tensor_index_of_layer(layer) + num_k_tensors. + assert self._layer_to_k_index == [ + None if t is None else v_index[id(t)] for t in device_kv_pool.v_buffer + ], "flat host mirror: K/V dedup orders diverge" + + state_slabs = _state_slabs(device_kv_pool) + state_tensors = [t for pair in state_slabs for t in pair] + + # layer -> slab pair index for state layers (identity-matched via + # the pool's occurrence-indexed get_state_buffers binding). + self._layer_to_state_pair: dict[int, int] = {} + if state_slabs: + pair_of_conv = {id(conv): n for n, (conv, _) in enumerate(state_slabs)} + for layer_id in range(len(device_kv_pool.k_buffer)): + try: + conv, _ssm = device_kv_pool.get_state_buffers(layer_id) + except ValueError: + continue # not a state layer + self._layer_to_state_pair[layer_id] = pair_of_conv[id(conv)] + + pin = torch.cuda.is_available() + kv_pairs = [ + ( + dev, + torch.zeros( + (self.num_host_pages * self.page_size, *dev.shape[1:]), + dtype=dev.dtype, + pin_memory=pin, + ), + ) + for dev in k_tensors + v_tensors + ] + state_pairs = [ + ( + dev, + torch.zeros( + (self.num_host_pages, *dev.shape[1:]), + dtype=dev.dtype, + pin_memory=pin, + ), + ) + for dev in state_tensors + ] + self.tensor_pairs: tuple[tuple[torch.Tensor, torch.Tensor], ...] = tuple( + kv_pairs + state_pairs + ) + # Rows one page spans on each pair: page_size token rows for KV, + # one page-indexed snapshot row for state slabs. + self.row_spans: tuple[int, ...] = (self.page_size,) * len(kv_pairs) + ( + 1, + ) * len(state_pairs) + + def tensor_index_of_layer(self, layer_id: int) -> int: + """Index of layer_id's K tensor in tensor_pairs (paired slab layers + share the index); its V tensor is at index + num_k_tensors. + Raises ValueError for flat GDN state layers (no KV tensor); fence + those on state_tensor_indices_of_layer instead.""" + index = self._layer_to_k_index[layer_id] + if index is None: + raise ValueError(f"layer {layer_id} is a state layer; it has no KV mirror") + return index + + def state_tensor_indices_of_layer(self, layer_id: int) -> tuple[int, int] | None: + """(conv_idx, ssm_idx) of layer_id's state slab pair in tensor_pairs + (conv immediately precedes its ssm), or None for layers without + state.""" + pair = self._layer_to_state_pair.get(layer_id) + if pair is None: + return None + base = 2 * self.num_k_tensors + 2 * pair + return base, base + 1 + + def bytes_per_host_page(self) -> int: + return sum( + dev.element_size() * dev[0].numel() * span + for (dev, _), span in zip(self.tensor_pairs, self.row_spans) + ) + + def _copy_pages( + self, + pairs: Iterable[tuple[int, int]], + stream, + to_host: bool, + record_events: bool, + ) -> list[torch.cuda.Event]: + pairs = list(pairs) + events: list[torch.cuda.Event] = [] + with torch.cuda.stream(stream): + for (dev, mirror), p in zip(self.tensor_pairs, self.row_spans): + for device_page, host_page in pairs: + dev_rows = dev[device_page * p : (device_page + 1) * p] + host_rows = mirror[host_page * p : (host_page + 1) * p] + if to_host: + host_rows.copy_(dev_rows, non_blocking=True) + else: + dev_rows.copy_(host_rows, non_blocking=True) + if record_events: + event = torch.cuda.Event() + event.record() + events.append(event) + return events + + def store_pages(self, pairs: Iterable[tuple[int, int]], stream) -> None: + """Copy each (device_page, host_page) pair device -> host on stream.""" + self._copy_pages(pairs, stream, to_host=True, record_events=False) + + def load_pages(self, pairs: Iterable[tuple[int, int]], stream) -> None: + """Copy each (device_page, host_page) pair host -> device on stream.""" + self._copy_pages(pairs, stream, to_host=False, record_events=False) + + def load_pages_with_events( + self, pairs: Iterable[tuple[int, int]], stream + ) -> list[torch.cuda.Event]: + """load_pages, recording one event per device tensor (tensor_pairs + order) after that tensor's copies -- D2's per-slab fencing hook.""" + return self._copy_pages(pairs, stream, to_host=False, record_events=True) diff --git a/python/tokenspeed/runtime/cache/kv_cache_host.py b/python/tokenspeed/runtime/cache/kv_cache_host.py new file mode 100755 index 0000000..7b0537c --- /dev/null +++ b/python/tokenspeed/runtime/cache/kv_cache_host.py @@ -0,0 +1,976 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the FluentLLM project +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import abc +import threading +from functools import wraps +from pathlib import Path + +import psutil +import torch +from tokenspeed_kernel.platform import current_platform + +from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool +from tokenspeed.runtime.layers.attention.kv_cache.dsa import DSATokenToKVPool +from tokenspeed.runtime.layers.attention.kv_cache.mha import MHATokenToKVPool +from tokenspeed.runtime.layers.attention.kv_cache.mla import MLATokenToKVPool +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) +_platform = current_platform() +if _platform.is_nvidia: + from tokenspeed_kernel.ops.kvcache.cuda import ( + transfer_kv_all_layer_lf_pf, + transfer_kv_all_layer_lf_ph, + transfer_kv_all_layer_mla, + transfer_kv_all_layer_mla_lf_pf, + transfer_kv_direct, + transfer_kv_per_layer_mla, + transfer_kv_per_layer_mla_pf_lf, + transfer_kv_per_layer_pf_lf, + transfer_kv_per_layer_ph_lf, + ) + +from tokenspeed_kernel.ops.kvcache.triton import ( + transfer_kv_all_layer, + transfer_kv_per_layer, +) + +if _platform.is_amd: + from tokenspeed_kernel.ops.kvcache.triton import ( + transfer_kv_all_layer_mla, + transfer_kv_per_layer_mla, + ) + +MLA_KVSTORE_LOADBACK_BLOCK_QUOTA = 16 +MLA_KVSTORE_WRITEBACK_BLOCK_QUOTA = 16 + + +def _read_cgroup_memory_value(path: Path) -> int | None: + try: + raw = path.read_text().strip() + except OSError: + return None + if not raw or raw == "max": + return None + try: + value = int(raw) + except ValueError: + return None + if value <= 0 or value >= (1 << 60): + return None + return value + + +def get_cgroup_memory_limit_and_current() -> tuple[int, int] | None: + """Return the active cgroup memory limit/current bytes, if constrained.""" + + limit = _read_cgroup_memory_value(Path("/sys/fs/cgroup/memory.max")) + current = _read_cgroup_memory_value(Path("/sys/fs/cgroup/memory.current")) + if limit is None or current is None: + limit = _read_cgroup_memory_value( + Path("/sys/fs/cgroup/memory/memory.limit_in_bytes") + ) + current = _read_cgroup_memory_value( + Path("/sys/fs/cgroup/memory/memory.usage_in_bytes") + ) + if limit is None or current is None: + return None + host_total = psutil.virtual_memory().total + if limit >= host_total: + return None + return limit, current + + +def get_available_host_memory_bytes( + reserve_bytes: int, +) -> tuple[int, int, int | None]: + host_available = max(psutil.virtual_memory().available - reserve_bytes, 0) + cgroup_available = None + cgroup_info = get_cgroup_memory_limit_and_current() + if cgroup_info is not None: + limit, current = cgroup_info + cgroup_available = max(limit - current - reserve_bytes, 0) + return min(host_available, cgroup_available), host_available, cgroup_available + return host_available, host_available, None + + +def synchronized(func): + @wraps(func) + def wrapper(self, *args, **kwargs): + with self.lock: + return func(self, *args, **kwargs) + + return wrapper + + +class HostKVCache(abc.ABC): + def __init__( + self, + device_pool: BaseTokenToKVPool, + host_to_device_ratio: float, + host_size: int, + page_size: int, + layout: str, + device: str, + host_size_tokens: int = 0, + ): + self.device_pool = device_pool + self.page_size = page_size + self.layout = layout + self.device = device + + self.dtype = device_pool.store_dtype + self.size_per_token = self.get_size_per_token() + if host_size_tokens > 0: + # Explicitly specified token count takes the highest priority. + # Used when this pool must share the same page address space as + # another host pool (e.g. draft model sharing base model pages). + self.size = host_size_tokens + elif host_size > 0: + self.size = int(host_size * 1e9 // self.size_per_token) + else: + self.size = int(device_pool.size * host_to_device_ratio) + # Align up the host memory pool size to the page size + self.page_num = self.size // self.page_size + 1 + self.size = self.page_num * self.page_size + + if self.size > device_pool.size: + logger.warning( + "The host memory is less than the device memory with the current protocol" + ) + + # Verify there is enough available host memory. + requested_bytes = self.size * self.size_per_token + # preserve at least 10GB for other usage + ten_gb = 10 * (1024**3) + available_bytes, host_available, cgroup_available = ( + get_available_host_memory_bytes(ten_gb) + ) + if requested_bytes > available_bytes: + raise ValueError( + f"Not enough host memory available. Requesting " + f"{requested_bytes / 1e9:.2f} GB but only have " + f"{available_bytes / 1e9:.2f} GB free. Please reduce the " + f"size of the KVStore." + ) + else: + logger.info( + "Allocating %.2f GB host memory for KVStore. host_size=%r self.size_per_token=%r host_to_device_ratio=%r device_pool.size=%r host_mem.available=%r", + requested_bytes / 1e9, + host_size, + self.size_per_token, + host_to_device_ratio, + device_pool.size, + host_available, + ) + if cgroup_available is not None: + logger.info( + "KVStore cgroup-aware available host memory: %.2f GB", + cgroup_available / 1e9, + ) + + self.kv_buffer = self.init_kv_buffer() + + # A lock for synchronized operations on memory allocation and state transitions. + self.lock = threading.RLock() + self.clear() + + @abc.abstractmethod + def get_size_per_token(self): + raise NotImplementedError() + + @abc.abstractmethod + def init_kv_buffer(self): + raise NotImplementedError() + + @abc.abstractmethod + def load_to_device_per_layer( + self, device_pool, host_indices, device_indices, layer_id, io_backend + ) -> None: + """ + Load KV data from the host memory pool to the device memory pool for a specific layer. + """ + raise NotImplementedError() + + @abc.abstractmethod + def backup_from_device_all_layer( + self, + device_pool, + host_indices, + device_indices, + io_backend, + block_quota: int | None = None, + ) -> None: + """ + Backup KV data from the device memory pool to the host memory pool for all layers. + """ + raise NotImplementedError() + + @abc.abstractmethod + def get_data_page(self, index, flat: bool = True) -> torch.Tensor: + """ + Get a flat data page from the host memory pool. + """ + raise NotImplementedError() + + @abc.abstractmethod + def get_dummy_flat_data_page(self) -> torch.Tensor: + """ + Get a dummy flat data page from the host memory pool. + This is used for prefetching or initializing empty pages. + """ + raise NotImplementedError() + + @abc.abstractmethod + def set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None: + """ + Set a flat data page to the host memory pool. + """ + raise NotImplementedError() + + @synchronized + def clear(self): + # Initialize memory states and tracking structures. + self.mem_state = torch.zeros( + (self.size,), dtype=torch.uint8, device=self.device + ) + self.free_slots = torch.arange(self.size, dtype=torch.int64) + + def available_size(self): + return len(self.free_slots) + + @synchronized + def alloc(self, need_size: int) -> torch.Tensor | None: + if need_size % self.page_size != 0: + raise ValueError("The requested size should be a multiple of page_size.") + if need_size > self.available_size(): + return None + + select_index = self.free_slots[:need_size] + self.free_slots = self.free_slots[need_size:] + + return select_index + + @synchronized + def free(self, indices: torch.Tensor) -> int: + self.free_slots = torch.cat([self.free_slots, indices]) + return len(indices) + + +class MHATokenToKVPoolHost(HostKVCache): + device_pool: MHATokenToKVPool + + def __init__( + self, + device_pool: MHATokenToKVPool, + host_to_device_ratio: float, + host_size: int, + page_size: int, + layout: str, + device: str = "cpu", + host_size_tokens: int = 0, + ): + super().__init__( + device_pool, + host_to_device_ratio, + host_size, + page_size, + layout, + device, + host_size_tokens=host_size_tokens, + ) + self.element_dim = self.device_pool.head_num * self.device_pool.head_dim + self.k_data_refs = [self.k_buffer[i] for i in range(self.layer_num)] + self.v_data_refs = [self.v_buffer[i] for i in range(self.layer_num)] + platform = current_platform() + self.k_data_ptrs = torch.tensor( + [platform.device_visible_data_ptr(x) for x in self.k_data_refs], + dtype=torch.uint64, + device=self.device_pool.device, + ) + self.v_data_ptrs = torch.tensor( + [platform.device_visible_data_ptr(x) for x in self.v_data_refs], + dtype=torch.uint64, + device=self.device_pool.device, + ) + + def get_size_per_token(self): + self.head_num = self.device_pool.head_num + self.head_dim = self.device_pool.head_dim + self.layer_num = self.device_pool.layer_num + + return self.head_dim * self.head_num * self.layer_num * self.dtype.itemsize * 2 + + def get_ksize_per_token(self): + return self.get_size_per_token() // 2 + + def init_kv_buffer(self): + if self.layout == "layer_first": + dims = (2, self.layer_num, self.size, self.head_num, self.head_dim) + elif self.layout == "page_first": + dims = (2, self.size, self.layer_num, self.head_num, self.head_dim) + elif self.layout == "page_head": + dims = ( + 2, + self.page_num, + self.head_num, + self.page_size, + self.layer_num, + self.head_dim, + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + self.token_stride_size = self.head_num * self.head_dim * self.dtype.itemsize + self.layout_dim = self.token_stride_size * self.layer_num + buffer = torch.empty( + dims, + dtype=self.dtype, + device=self.device, + ) + current_platform().register_host_tensor_for_gpu_access(buffer) + return buffer + + @property + def k_buffer(self): + return self.kv_buffer[0] + + @property + def v_buffer(self): + return self.kv_buffer[1] + + def load_to_device_per_layer( + self, + device_pool, + host_indices, + device_indices, + layer_id, + io_backend, + ): + if io_backend == "kernel": + if self.layout == "layer_first": + transfer_kv_per_layer( + src_k=self.k_buffer[layer_id], + dst_k=device_pool.k_buffer[layer_id], + src_v=self.v_buffer[layer_id], + dst_v=device_pool.v_buffer[layer_id], + src_indices=host_indices, + dst_indices=device_indices, + item_size=self.token_stride_size, + ) + elif self.layout == "page_first": + transfer_kv_per_layer_pf_lf( + src_k=self.k_buffer, + dst_k=device_pool.k_buffer[layer_id], + src_v=self.v_buffer, + dst_v=device_pool.v_buffer[layer_id], + src_indices=host_indices, + dst_indices=device_indices, + layer_id=layer_id, + item_size=self.token_stride_size, + src_layout_dim=self.layout_dim, + ) + elif self.layout == "page_head": + transfer_kv_per_layer_ph_lf( + src_k=self.k_buffer, + dst_k=device_pool.k_buffer[layer_id], + src_v=self.v_buffer, + dst_v=device_pool.v_buffer[layer_id], + src_indices=host_indices, + dst_indices=device_indices, + layer_id=layer_id, + item_size=self.token_stride_size, + src_layout_dim=self.layout_dim, + page_size=self.page_size, + head_num=self.head_num, + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + elif io_backend == "direct": + if self.layout == "layer_first": + transfer_kv_direct( + src_layers=[self.k_buffer[layer_id], self.v_buffer[layer_id]], + dst_layers=[ + device_pool.k_buffer[layer_id], + device_pool.v_buffer[layer_id], + ], + src_indices=host_indices, + dst_indices=device_indices, + page_size=self.page_size, + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + else: + raise ValueError(f"Unsupported IO backend: {io_backend}") + + def backup_from_device_all_layer( + self, + device_pool, + host_indices, + device_indices, + io_backend, + block_quota: int | None = None, + ): + if io_backend == "kernel": + if self.layout == "layer_first": + transfer_kv_all_layer( + src_k_layers=device_pool.k_data_ptrs, + dst_k_layers=self.k_data_ptrs, + src_v_layers=device_pool.v_data_ptrs, + dst_v_layers=self.v_data_ptrs, + src_indices=device_indices, + dst_indices=host_indices, + item_size=self.token_stride_size, + num_layers=self.layer_num, + ) + elif self.layout == "page_first": + transfer_kv_all_layer_lf_pf( + src_k_layers=device_pool.k_data_ptrs, + dst_k=self.k_buffer, + src_v_layers=device_pool.v_data_ptrs, + dst_v=self.v_buffer, + src_indices=device_indices, + dst_indices=host_indices, + item_size=self.token_stride_size, + dst_layout_dim=self.layout_dim, + num_layers=self.layer_num, + ) + elif self.layout == "page_head": + transfer_kv_all_layer_lf_ph( + src_k_layers=device_pool.k_data_ptrs, + dst_k=self.k_buffer, + src_v_layers=device_pool.v_data_ptrs, + dst_v=self.v_buffer, + src_indices=device_indices, + dst_indices=host_indices, + item_size=self.token_stride_size, + dst_layout_dim=self.layout_dim, + num_layers=self.layer_num, + page_size=self.page_size, + head_num=self.head_num, + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + elif io_backend == "direct": + if self.layout == "layer_first": + transfer_kv_direct( + src_layers=device_pool.k_buffer + device_pool.v_buffer, + dst_layers=self.k_data_refs + self.v_data_refs, + src_indices=device_indices, + dst_indices=host_indices, + page_size=self.page_size, + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + else: + raise ValueError(f"Unsupported IO backend: {io_backend}") + + def get_data_page(self, index, flat: bool = True) -> torch.Tensor: + if self.layout == "layer_first": + data_page = self.kv_buffer[:, :, index : index + self.page_size, :, :] + elif self.layout == "page_first": + data_page = self.kv_buffer[:, index : index + self.page_size, :, :, :] + elif self.layout == "page_head": + real_index = index // self.page_size + data_page = self.kv_buffer[:, real_index : real_index + 1, :, :, :, :] + else: + raise ValueError(f"Unsupported layout: {self.layout}") + if flat: + data_page = data_page.flatten() + return data_page + + def get_dummy_flat_data_page(self) -> torch.Tensor: + return torch.zeros( + (2, self.layer_num, self.page_size, self.head_num, self.head_dim), + dtype=self.dtype, + device=self.device, + pin_memory=True, + ).flatten() + + def set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None: + if self.layout == "layer_first": + self.kv_buffer[:, :, index : index + self.page_size, :, :] = ( + data_page.reshape( + 2, + self.layer_num, + self.page_size, + self.head_num, + self.head_dim, + ) + ) + elif self.layout == "page_first": + self.kv_buffer[:, index : index + self.page_size, :, :, :] = ( + data_page.reshape( + 2, self.page_size, self.layer_num, self.head_num, self.head_dim + ) + ) + elif self.layout == "page_head": + real_index = index // self.page_size + self.kv_buffer[:, real_index : real_index + 1, :, :, :, :] = ( + data_page.reshape( + 2, 1, self.head_num, self.page_size, self.layer_num, self.head_dim + ) + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + + def get_page_buffer_meta(self, indices): + if len(indices) % self.page_size != 0: + raise ValueError("indices length must be a multiple of page_size") + ptr_list = [] + kv_buffer_data_ptr = self.kv_buffer.data_ptr() + indices = indices.tolist() + v_offset = ( + self.layer_num + * self.size + * self.head_num + * self.head_dim + * self.dtype.itemsize + ) + if self.layout == "layer_first": + for index in range(0, len(indices), self.page_size): + for layer_id in range(self.layer_num): + k_ptr = ( + kv_buffer_data_ptr + + indices[index] + * self.head_num + * self.head_dim + * self.dtype.itemsize + + layer_id + * self.size + * self.head_num + * self.head_dim + * self.dtype.itemsize + ) + v_ptr = k_ptr + v_offset + ptr_list.append(k_ptr) + ptr_list.append(v_ptr) + element_size = ( + self.dtype.itemsize * self.page_size * self.head_num * self.head_dim + ) + element_size_list = [element_size] * len(ptr_list) + elif self.layout in ["page_first", "page_head"]: + for index in range(0, len(indices), self.page_size): + k_ptr = ( + kv_buffer_data_ptr + + indices[index] + * self.layer_num + * self.head_num + * self.head_dim + * self.dtype.itemsize + ) + v_ptr = k_ptr + v_offset + ptr_list.append(k_ptr) + ptr_list.append(v_ptr) + element_size = ( + self.layer_num + * self.dtype.itemsize + * self.page_size + * self.head_num + * self.head_dim + ) + element_size_list = [element_size] * len(ptr_list) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + return ptr_list, element_size_list + + +class MLATokenToKVPoolHost(HostKVCache): + device_pool: MLATokenToKVPool + + def __init__( + self, + device_pool: MLATokenToKVPool, + host_to_device_ratio: float, + host_size: int, + page_size: int, + layout: str, + device: str = "cpu", + host_size_tokens: int = 0, + ): + super().__init__( + device_pool, + host_to_device_ratio, + host_size, + page_size, + layout, + device, + host_size_tokens=host_size_tokens, + ) + self.data_refs = [self.kv_buffer[i] for i in range(self.layer_num)] + platform = current_platform() + self.data_ptrs = torch.tensor( + [platform.device_visible_data_ptr(x) for x in self.data_refs], + dtype=torch.uint64, + device=self.device_pool.device, + ) + + def get_size_per_token(self): + self.kv_lora_rank = self.device_pool.kv_lora_rank + self.qk_rope_head_dim = self.device_pool.qk_rope_head_dim + self.layer_num = self.device_pool.layer_num + + return ( + (self.kv_lora_rank + self.qk_rope_head_dim) + * 1 + * self.dtype.itemsize + * self.layer_num + ) + + def get_ksize_per_token(self): + return self.get_size_per_token() + + def init_kv_buffer(self): + if self.layout == "layer_first": + dims = ( + self.layer_num, + self.size, + 1, + self.kv_lora_rank + self.qk_rope_head_dim, + ) + elif self.layout == "page_first": + dims = ( + self.size, + self.layer_num, + 1, + self.kv_lora_rank + self.qk_rope_head_dim, + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + self.token_stride_size = ( + self.kv_lora_rank + self.qk_rope_head_dim + ) * self.dtype.itemsize + self.layout_dim = self.token_stride_size * self.layer_num + buffer = torch.empty( + dims, + dtype=self.dtype, + device=self.device, + ) + current_platform().register_host_tensor_for_gpu_access(buffer) + return buffer + + def load_to_device_per_layer( + self, device_pool, host_indices, device_indices, layer_id, io_backend + ): + if io_backend == "kernel": + if self.layout == "layer_first": + transfer_kv_per_layer_mla( + src=self.kv_buffer[layer_id], + dst=device_pool.kv_buffer[layer_id], + src_indices=host_indices, + dst_indices=device_indices, + item_size=self.token_stride_size, + block_quota=MLA_KVSTORE_LOADBACK_BLOCK_QUOTA, + ) + elif self.layout == "page_first": + transfer_kv_per_layer_mla_pf_lf( + src=self.kv_buffer, + dst=device_pool.kv_buffer[layer_id], + src_indices=host_indices, + dst_indices=device_indices, + layer_id=layer_id, + item_size=self.token_stride_size, + src_layout_dim=self.layout_dim, + block_quota=MLA_KVSTORE_LOADBACK_BLOCK_QUOTA, + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + elif io_backend == "direct": + if self.layout == "layer_first": + transfer_kv_direct( + src_layers=[self.kv_buffer[layer_id]], + dst_layers=[device_pool.kv_buffer[layer_id]], + src_indices=host_indices, + dst_indices=device_indices, + page_size=self.page_size, + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + else: + raise ValueError(f"Unsupported IO backend: {io_backend}") + + def backup_from_device_all_layer( + self, + device_pool, + host_indices, + device_indices, + io_backend, + block_quota: int | None = None, + ): + if block_quota is None: + block_quota = MLA_KVSTORE_WRITEBACK_BLOCK_QUOTA + if io_backend == "kernel": + if self.layout == "layer_first": + transfer_kv_all_layer_mla( + src_layers=device_pool.data_ptrs, + dst_layers=self.data_ptrs, + src_indices=device_indices, + dst_indices=host_indices, + item_size=self.token_stride_size, + num_layers=self.layer_num, + block_quota=block_quota, + ) + elif self.layout == "page_first": + transfer_kv_all_layer_mla_lf_pf( + src_layers=device_pool.data_ptrs, + dst=self.kv_buffer, + src_indices=device_indices, + dst_indices=host_indices, + item_size=self.token_stride_size, + dst_layout_dim=self.layout_dim, + num_layers=self.layer_num, + block_quota=block_quota, + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + elif io_backend == "direct": + if self.layout == "layer_first": + transfer_kv_direct( + src_layers=device_pool.kv_buffer, + dst_layers=self.data_refs, + src_indices=device_indices, + dst_indices=host_indices, + page_size=self.page_size, + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + else: + raise ValueError(f"Unsupported IO backend: {io_backend}") + + def get_data_page(self, index, flat: bool = True) -> torch.Tensor: + if self.layout == "layer_first": + data_page = self.kv_buffer[:, index : index + self.page_size, :, :] + elif self.layout == "page_first": + data_page = self.kv_buffer[index : index + self.page_size, :, :, :] + else: + raise ValueError(f"Unsupported layout: {self.layout}") + if flat: + data_page = data_page.flatten() + return data_page + + def get_dummy_flat_data_page(self) -> torch.Tensor: + return torch.zeros( + ( + self.layer_num, + self.page_size, + 1, + self.kv_lora_rank + self.qk_rope_head_dim, + ), + dtype=self.dtype, + device=self.device, + pin_memory=True, + ).flatten() + + def set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None: + if self.layout == "layer_first": + self.kv_buffer[:, index : index + self.page_size, :, :] = data_page.reshape( + self.layer_num, + self.page_size, + 1, + self.kv_lora_rank + self.qk_rope_head_dim, + ) + elif self.layout == "page_first": + self.kv_buffer[index : index + self.page_size, :, :, :] = data_page.reshape( + self.page_size, + self.layer_num, + 1, + self.kv_lora_rank + self.qk_rope_head_dim, + ) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + + def get_page_buffer_meta(self, indices): + if len(indices) % self.page_size != 0: + raise ValueError("indices length must be a multiple of page_size") + ptr_list = [] + kv_buffer_data_ptr = self.kv_buffer.data_ptr() + indices = indices.tolist() + if self.layout == "layer_first": + for index in range(0, len(indices), self.page_size): + for layer_id in range(self.layer_num): + k_ptr = ( + kv_buffer_data_ptr + + indices[index] + * (self.kv_lora_rank + self.qk_rope_head_dim) + * self.dtype.itemsize + + layer_id + * self.size + * (self.kv_lora_rank + self.qk_rope_head_dim) + * self.dtype.itemsize + ) + ptr_list.append(k_ptr) + element_size = ( + self.dtype.itemsize + * self.page_size + * (self.kv_lora_rank + self.qk_rope_head_dim) + ) + element_size_list = [element_size] * len(ptr_list) + elif self.layout == "page_first": + for index in range(0, len(indices), self.page_size): + k_ptr = ( + kv_buffer_data_ptr + + indices[index] + * self.layer_num + * (self.kv_lora_rank + self.qk_rope_head_dim) + * self.dtype.itemsize + ) + ptr_list.append(k_ptr) + element_size = ( + self.layer_num + * self.dtype.itemsize + * self.page_size + * (self.kv_lora_rank + self.qk_rope_head_dim) + ) + element_size_list = [element_size] * len(ptr_list) + else: + raise ValueError(f"Unsupported layout: {self.layout}") + return ptr_list, element_size_list + + +class DSATokenToKVPoolHost(MLATokenToKVPoolHost): + """Host (L2) mirror of the GLM DSA KV pool. + + Extends the MLA latent host pool with the DSA FP8 index-K buffer. Both buffers + mirror the device row layout and transfer per token. The index-K buffers are + in a block-split layout: each page is laid out as + ``[page_size * head_dim FP8 values]`` followed by ``[page_size * num_groups FP32 scales]``. + Hence, it requires the token indices are built as whole page-expanded blocks + (see host_executor.page_ids_to_token_indices); otherwise the D<->H transfer + would be corrupted. + """ + + device_pool: DSATokenToKVPool + + def __init__( + self, + device_pool: DSATokenToKVPool, + host_to_device_ratio: float, + host_size: int, + page_size: int, + layout: str, + device: str = "cpu", + host_size_tokens: int = 0, + ): + if device_pool.quant_method == "per_token_head": + raise NotImplementedError( + "DSA KVStore does not support the per_token_head latent layout." + ) + if layout != "layer_first": + raise NotImplementedError( + f"DSA KVStore supports only the layer_first host layout, got {layout}." + ) + self.index_k_row_bytes = device_pool.index_k_row_bytes + super().__init__( + device_pool, + host_to_device_ratio, + host_size, + page_size, + layout, + device, + host_size_tokens=host_size_tokens, + ) + self.index_k_refs = [self.index_k_buffer[i] for i in range(self.layer_num)] + platform = current_platform() + self.index_k_data_ptrs = torch.tensor( + [platform.device_visible_data_ptr(x) for x in self.index_k_refs], + dtype=torch.uint64, + device=self.device_pool.device, + ) + + def get_size_per_token(self): + return super().get_size_per_token() + self.index_k_row_bytes * self.layer_num + + def init_kv_buffer(self): + kv_buffer = super().init_kv_buffer() + # Mirror the device index-K layout: page p of layer L occupies rows + # [p * page_size : (p + 1) * page_size], so a whole page is contiguous + # and the block-split FP8/scale bytes within it survive a raw page copy. + + self.index_k_buffer = torch.zeros( + (self.layer_num, self.size, self.index_k_row_bytes), + dtype=torch.uint8, + device=self.device, + ) + current_platform().register_host_tensor_for_gpu_access(self.index_k_buffer) + return kv_buffer + + def load_to_device_per_layer( + self, device_pool, host_indices, device_indices, layer_id, io_backend + ): + super().load_to_device_per_layer( + device_pool, host_indices, device_indices, layer_id, io_backend + ) + if io_backend == "kernel": + transfer_kv_per_layer_mla( + src=self.index_k_buffer[layer_id], + dst=device_pool.index_k_buffer[layer_id], + src_indices=host_indices, + dst_indices=device_indices, + item_size=self.index_k_row_bytes, + block_quota=MLA_KVSTORE_LOADBACK_BLOCK_QUOTA, + ) + elif io_backend == "direct": + transfer_kv_direct( + src_layers=[self.index_k_buffer[layer_id]], + dst_layers=[device_pool.index_k_buffer[layer_id]], + src_indices=host_indices, + dst_indices=device_indices, + page_size=self.page_size, + ) + else: + raise ValueError(f"Unsupported IO backend: {io_backend}") + + def backup_from_device_all_layer( + self, + device_pool, + host_indices, + device_indices, + io_backend, + block_quota: int | None = None, + ): + super().backup_from_device_all_layer( + device_pool, host_indices, device_indices, io_backend, block_quota + ) + if io_backend == "kernel": + if block_quota is None: + block_quota = MLA_KVSTORE_WRITEBACK_BLOCK_QUOTA + transfer_kv_all_layer_mla( + src_layers=device_pool.index_k_data_ptrs, + dst_layers=self.index_k_data_ptrs, + src_indices=device_indices, + dst_indices=host_indices, + item_size=self.index_k_row_bytes, + num_layers=self.layer_num, + block_quota=block_quota, + ) + elif io_backend == "direct": + transfer_kv_direct( + src_layers=list(device_pool.index_k_buffer), + dst_layers=self.index_k_refs, + src_indices=device_indices, + dst_indices=host_indices, + page_size=self.page_size, + ) + else: + raise ValueError(f"Unsupported IO backend: {io_backend}") diff --git a/python/tokenspeed/runtime/cache/kvstore_controller.py b/python/tokenspeed/runtime/cache/kvstore_controller.py new file mode 100755 index 0000000..176b73b --- /dev/null +++ b/python/tokenspeed/runtime/cache/kvstore_controller.py @@ -0,0 +1,85 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from collections.abc import Iterable + +from tokenspeed.runtime.utils import get_colorful_logger, get_device_module + +logger = get_colorful_logger(__name__) + +device_module = get_device_module() + + +class LayerLoadingEvent: + def __init__(self, num_layers: int): + self._num_layers = num_layers + self.load_events = [device_module.Event() for _ in range(num_layers)] + self.start_event = device_module.Event() # start event on controller stream + + def complete(self, layer_index: int) -> None: + if not 0 <= layer_index < self._num_layers: + raise IndexError(f"layer_index out of range: {layer_index}") + self.load_events[layer_index].record() + + def wait(self, layer_index: int) -> None: + device_module.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_indices: tuple[int, ...] = () + + def update_producer(self): + next_index = (self.producer_index + 1) % self.num_counters + if not self.events[next_index].finish_event.query(): + self.events[next_index].finish_event.synchronize() + self.producer_index = next_index + return self.producer_index + + def set_consumer(self, indices: int | Iterable[int]): + if isinstance(indices, int): + self.consumer_indices = () if indices < 0 else (indices,) + return + deduped = [] + for index in indices: + if index >= 0 and index not in deduped: + deduped.append(index) + self.consumer_indices = tuple(deduped) + + def wait_until(self, threshold: int): + if not self.consumer_indices: + return + for consumer_index in self.consumer_indices: + self.events[consumer_index].wait(threshold) + + def reset(self): + self.producer_index = -1 + self.consumer_indices = () diff --git a/python/tokenspeed/runtime/cache/kvstore_storage.py b/python/tokenspeed/runtime/cache/kvstore_storage.py new file mode 100755 index 0000000..98360db --- /dev/null +++ b/python/tokenspeed/runtime/cache/kvstore_storage.py @@ -0,0 +1,158 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any + +import torch + +from tokenspeed.runtime.cache.kv_cache_host import HostKVCache +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +@dataclass +class KVStoreStorageConfig: + tp_rank: int + tp_size: int + is_mla_model: bool + is_page_first_layout: bool + model_name: str | None + extra_config: dict | None = None + + +@dataclass +class KVStoreStorageExtraInfo: + prefix_keys: list[str] | None = None + extra_info: dict | None = None + + +class KVStoreStorage(ABC): + """ + KVStoreStorage is a class that provides a generic key-value interface for storing and retrieving KV cache. + It abstracts the underlying storage mechanism, allowing different implementations to be used. + """ + + def register_mem_pool_host(self, mem_pool_host: HostKVCache) -> None: + self.mem_pool_host = mem_pool_host + + def batch_get_v1( + self, + keys: list[str], + host_indices: torch.Tensor, + extra_info: KVStoreStorageExtraInfo | None = None, + ) -> list[bool]: + """ + Retrieve values for multiple keys. + Returns a list of booleans indicating success for each key. + """ + raise NotImplementedError + + def batch_set_v1( + self, + keys: list[str], + host_indices: torch.Tensor, + extra_info: KVStoreStorageExtraInfo | None = None, + ) -> list[bool]: + """ + Store multiple key-value pairs. + Returns a list of booleans indicating success for each key. + """ + raise NotImplementedError + + @abstractmethod + def get( + self, + key: str, + target_location: Any | None = None, + target_sizes: Any | None = None, + ) -> torch.Tensor | None: + """ + Retrieve the value associated with the given key. + Returns None if the key does not exist. + """ + raise NotImplementedError + + @abstractmethod + def batch_get( + self, + keys: list[str], + target_locations: Any | None = None, + target_sizes: Any | None = None, + ) -> list[torch.Tensor | None] | int: + """ + Retrieve values for multiple keys. + Returns a list of tensors or None for each key. + """ + raise NotImplementedError + + @abstractmethod + def set( + self, + key: str, + value: Any | None = None, + target_location: Any | None = None, + target_sizes: Any | None = None, + ) -> bool: + """ + Store the value associated with the given key. + Returns True if the operation was successful, False otherwise. + """ + raise NotImplementedError + + @abstractmethod + def batch_set( + self, + keys: list[str], + values: Any | None = None, + target_locations: Any | None = None, + target_sizes: Any | None = None, + ) -> bool: + """ + Store multiple key-value pairs. + Returns True if all operations were successful, False otherwise. + """ + raise NotImplementedError + + @abstractmethod + def exists(self, key: str) -> bool: + """ + Check if the key exists in the storage. + Returns True if the key exists, False otherwise. + """ + raise NotImplementedError + + def batch_exists( + self, keys: list[str], extra_info: KVStoreStorageExtraInfo | None = None + ) -> int: + """ + Check if the keys exist in the storage. + return the number of consecutive existing keys from the start. + Can be overridden by subclasses for more efficient implementation. + """ + for index, key in enumerate(keys): + if not self.exists(key): + return index + return len(keys) + + def clear(self) -> None: + return None diff --git a/python/tokenspeed/runtime/cache/mamba_cache_host.py b/python/tokenspeed/runtime/cache/mamba_cache_host.py new file mode 100644 index 0000000..e19a39e --- /dev/null +++ b/python/tokenspeed/runtime/cache/mamba_cache_host.py @@ -0,0 +1,279 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import threading +from functools import wraps + +import torch +from tokenspeed_kernel.ops.kvcache.cuda import ( + transfer_kv_all_layer_mla, + transfer_kv_direct, + transfer_kv_per_layer_mla, +) +from tokenspeed_kernel.platform import current_platform + +from tokenspeed.runtime.layers.attention.backends.hybrid_linear_attn import ( + SimpleMambaPool, +) +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) +MAMBA_KVSTORE_LOADBACK_BLOCK_QUOTA = 16 +MAMBA_KVSTORE_WRITEBACK_BLOCK_QUOTA = 16 + + +def synchronized(func): + @wraps(func) + def wrapper(self, *args, **kwargs): + with self.lock: + return func(self, *args, **kwargs) + + return wrapper + + +class MambaPoolHost: + """Pinned host mirror for SimpleMambaPool conv_state and ssm_state.""" + + def __init__( + self, + device_pool: SimpleMambaPool, + host_size_slots: int, + layout: str = "layer_first", + pin_memory: bool = True, + device: str = "cpu", + register_host: bool = True, + ): + if layout != "layer_first": + raise ValueError("MambaPoolHost v1 only supports layer_first layout") + if host_size_slots <= 0: + raise ValueError("host_size_slots must be positive") + self.device_pool = device_pool + self.layout = layout + self.device = device + self.size = int(host_size_slots) + self.page_size = 1 + self.num_layers = int(device_pool.conv_state.shape[0]) + self.conv_shape = tuple(device_pool.conv_state.shape[2:]) + self.ssm_shape = tuple(device_pool.ssm_state.shape[2:]) + self.conv_dtype = device_pool.conv_state.dtype + self.ssm_dtype = device_pool.ssm_state.dtype + self.conv_item_size = device_pool.conv_state[0, 0].nbytes + self.ssm_item_size = device_pool.ssm_state[0, 0].nbytes + self.size_per_slot = self.num_layers * ( + self.conv_item_size + self.ssm_item_size + ) + + # cudaHostRegister pins ordinary host memory for GPU-side access. + # Avoid allocating an already pinned tensor when we will register it, + # because some CUDA stacks reject double registration. + use_pin_memory = bool(pin_memory and device == "cpu" and not register_host) + self.conv_buffer = torch.empty( + (self.num_layers, self.size, *self.conv_shape), + dtype=self.conv_dtype, + device=device, + pin_memory=use_pin_memory, + ) + self.ssm_buffer = torch.empty( + (self.num_layers, self.size, *self.ssm_shape), + dtype=self.ssm_dtype, + device=device, + pin_memory=use_pin_memory, + ) + if register_host: + platform = current_platform() + platform.register_host_tensor_for_gpu_access(self.conv_buffer) + platform.register_host_tensor_for_gpu_access(self.ssm_buffer) + + self.conv_data_refs = [self.conv_buffer[i] for i in range(self.num_layers)] + self.ssm_data_refs = [self.ssm_buffer[i] for i in range(self.num_layers)] + # Keep CUDA all-layer kernel pointer tables alive across async launches. + self._kernel_ptr_tables: dict[str, torch.Tensor] | None = None + + self.lock = threading.RLock() + self.clear() + logger.info( + "[mamba_l2] alloc host buffer pool_type=%s size_slots=%s " + "size_per_slot_mb=%.2f num_mamba_layers=%s layout=%s pin_memory=%s " + "total_gb=%.2f", + type(self).__name__, + self.size, + self.size_per_slot / 1e6, + self.num_layers, + self.layout, + use_pin_memory, + self.size * self.size_per_slot / 1e9, + ) + + @synchronized + def clear(self) -> None: + self.free_slots = torch.arange(self.size, dtype=torch.int64) + + def available_size(self) -> int: + return len(self.free_slots) + + @synchronized + def alloc(self, need_size: int) -> torch.Tensor | None: + if need_size <= 0: + return torch.empty((0,), dtype=torch.int64) + if need_size > self.available_size(): + logger.warning( + "[mamba_l2] alloc FAILED n=%s remain=%s (will trigger eviction)", + need_size, + self.available_size(), + ) + return None + selected = self.free_slots[:need_size] + self.free_slots = self.free_slots[need_size:] + logger.debug( + "[mamba_l2] alloc n=%s remain=%s", need_size, self.available_size() + ) + return selected + + @synchronized + def free(self, indices: torch.Tensor) -> int: + indices = indices.to(dtype=torch.int64, device="cpu") + self.free_slots = torch.cat([self.free_slots, indices]) + logger.debug( + "[mamba_l2] free n=%s deferred=%s remain=%s", + len(indices), + False, + self.available_size(), + ) + return len(indices) + + def backup_from_device_all_layer( + self, + device_pool: SimpleMambaPool, + host_indices: torch.Tensor, + device_indices: torch.Tensor, + io_backend: str, + block_quota: int | None = None, + ) -> None: + if block_quota is None: + block_quota = MAMBA_KVSTORE_WRITEBACK_BLOCK_QUOTA + if io_backend == "kernel": + ptrs = self._ensure_kernel_ptr_tables(device_pool) + transfer_kv_all_layer_mla( + src_layers=ptrs["device_conv"], + dst_layers=ptrs["host_conv"], + src_indices=device_indices, + dst_indices=host_indices, + item_size=self.conv_item_size, + num_layers=self.num_layers, + block_quota=block_quota, + ) + transfer_kv_all_layer_mla( + src_layers=ptrs["device_ssm"], + dst_layers=ptrs["host_ssm"], + src_indices=device_indices, + dst_indices=host_indices, + item_size=self.ssm_item_size, + num_layers=self.num_layers, + block_quota=block_quota, + ) + elif io_backend == "direct": + transfer_kv_direct( + src_layers=self._layer_refs(device_pool.conv_state) + + self._layer_refs(device_pool.ssm_state), + dst_layers=self.conv_data_refs + self.ssm_data_refs, + src_indices=device_indices, + dst_indices=host_indices, + page_size=self.page_size, + ) + else: + raise ValueError(f"Unsupported IO backend: {io_backend}") + + def load_to_device_per_layer( + self, + device_pool: SimpleMambaPool, + host_indices: torch.Tensor, + device_indices: torch.Tensor, + layer_idx: int, + io_backend: str = "kernel", + ) -> None: + if not 0 <= layer_idx < self.num_layers: + raise IndexError(f"layer_idx out of range: {layer_idx}") + if io_backend == "kernel": + transfer_kv_per_layer_mla( + src=self.conv_buffer[layer_idx], + dst=device_pool.conv_state[layer_idx], + src_indices=host_indices, + dst_indices=device_indices, + item_size=self.conv_item_size, + block_quota=MAMBA_KVSTORE_LOADBACK_BLOCK_QUOTA, + ) + transfer_kv_per_layer_mla( + src=self.ssm_buffer[layer_idx], + dst=device_pool.ssm_state[layer_idx], + src_indices=host_indices, + dst_indices=device_indices, + item_size=self.ssm_item_size, + block_quota=MAMBA_KVSTORE_LOADBACK_BLOCK_QUOTA, + ) + elif io_backend == "direct": + transfer_kv_direct( + src_layers=[self.conv_buffer[layer_idx], self.ssm_buffer[layer_idx]], + dst_layers=[ + device_pool.conv_state[layer_idx], + device_pool.ssm_state[layer_idx], + ], + src_indices=host_indices, + dst_indices=device_indices, + page_size=self.page_size, + ) + else: + raise ValueError(f"Unsupported IO backend: {io_backend}") + + def get_hybrid_pool_buffer(self) -> list[torch.Tensor]: + return [self.conv_buffer, self.ssm_buffer] + + @staticmethod + def _layer_refs(buffer: torch.Tensor) -> list[torch.Tensor]: + return [buffer[i] for i in range(buffer.shape[0])] + + def _ensure_kernel_ptr_tables( + self, device_pool: SimpleMambaPool + ) -> dict[str, torch.Tensor]: + if self._kernel_ptr_tables is None: + self._kernel_ptr_tables = { + "device_conv": self._data_ptrs( + device_pool.conv_state, device_pool.device + ), + "host_conv": self._data_ptrs(self.conv_buffer, device_pool.device), + "device_ssm": self._data_ptrs( + device_pool.ssm_state, device_pool.device + ), + "host_ssm": self._data_ptrs(self.ssm_buffer, device_pool.device), + } + return self._kernel_ptr_tables + + @staticmethod + def _data_ptrs(buffer: torch.Tensor, device) -> torch.Tensor: + platform = current_platform() + return torch.tensor( + [ + platform.device_visible_data_ptr(buffer[i]) + for i in range(buffer.shape[0]) + ], + dtype=torch.uint64, + device=device, + ) diff --git a/python/tokenspeed/runtime/cache/prefix_cache.py b/python/tokenspeed/runtime/cache/prefix_cache.py new file mode 100755 index 0000000..5fd5424 --- /dev/null +++ b/python/tokenspeed/runtime/cache/prefix_cache.py @@ -0,0 +1,647 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +""" +The prefix tree data structure for managing the KV cache. +""" + +from __future__ import annotations + +import dataclasses +import heapq +import time +from collections import defaultdict +from collections.abc import Callable +from typing import TYPE_CHECKING + +import torch + +from tokenspeed.runtime.cache.allocator import KVAllocator +from tokenspeed.runtime.cache.base_prefix_cache import BasePrefixCache, MatchResult +from tokenspeed.runtime.cache.evict_policy import ( + EvictionStrategy, + FIFOStrategy, + FILOStrategy, + LFUStrategy, + LRUStrategy, + MRUStrategy, + PriorityStrategy, +) +from tokenspeed.runtime.cache.req_to_token_pool import ReqToTokenPool +from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + +if TYPE_CHECKING: + from tokenspeed.runtime.engine.request import Req + +# Import KV events classes if needed +try: + from tokenspeed.runtime.pd.kv_events import ( + AllBlocksCleared, + BlockRemoved, + BlockStored, + ) + + KV_EVENTS_AVAILABLE = True +except ImportError: + BlockStored = None + BlockRemoved = None + AllBlocksCleared = None + KV_EVENTS_AVAILABLE = False + logger.warning( + "KV Events module not available. " + "KV cache events will not be emitted even if enable_kv_cache_events=True." + ) + + +@dataclasses.dataclass +class CacheInitParams: + disable: bool + req_to_token_pool: ReqToTokenPool + token_to_kv_pool_allocator: KVAllocator + page_size: int + + token_to_kv_pool: BaseTokenToKVPool = None + tp_cache_group: torch.distributed.ProcessGroup | None = None + eviction_policy: str = "lru" + disable_finished_insert: bool = False + + enable_metrics: bool = False + enable_kv_cache_events: bool = False + + +class TreeNode: + counter = 0 + + def __init__(self, id: int | None = None, priority: int = 0): + self.children = defaultdict(TreeNode) + self.parent = None + self.key = None + self.value = None + self.lock_ref = 0 + self.last_access_time = time.monotonic() + self.creation_time = time.monotonic() + + self.id = TreeNode.counter if id is None else id + TreeNode.counter += 1 + + ##### KVStore ##### + # host_value is token level, others are page level + self.hit_count = 0 + # indicating the node is locked to protect from eviction + # incremented when the node is referenced by a storage operation + self.host_ref_counter = 0 + # store the host indices of KV cache + self.host_value: torch.Tensor | None = None + # store hash values of each pages + self.hash_value: list[str] | None = None + # priority for priority-aware eviction + self.priority = priority + + @property + def evicted(self): + return self.value is None + + def __lt__(self, other: "TreeNode"): + return self.last_access_time < other.last_access_time + + +def _key_match(key0: list, key1: list): + i = 0 + for k0, k1 in zip(key0, key1): + if k0 != k1: + break + i += 1 + return i + + +class PrefixCache(BasePrefixCache): + def __init__(self, params: CacheInitParams): + self.disable = params.disable + self.req_to_token_pool = params.req_to_token_pool + self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator + self.enable_kv_cache_events = params.enable_kv_cache_events + self.page_size = ( + 1 + if params.token_to_kv_pool_allocator is None + else params.token_to_kv_pool_allocator.page_size + ) + self.eviction_policy = params.eviction_policy.lower() + self.key_match_fn = _key_match + + strategy_types: dict[str, type[EvictionStrategy]] = { + "lru": LRUStrategy, + "lfu": LFUStrategy, + "fifo": FIFOStrategy, + "mru": MRUStrategy, + "filo": FILOStrategy, + "priority": PriorityStrategy, + } + strategy_type = strategy_types.get(self.eviction_policy) + if strategy_type is None: + raise ValueError( + f"Unknown eviction policy: {self.eviction_policy}. Supported policies: 'lru', 'lfu', 'fifo', 'mru', 'filo', 'priority'." + ) + self.eviction_strategy = strategy_type() + # Initialize KV event queue + self.kv_event_queue = [] + + # Warn if KV events are enabled but module is not available + if self.enable_kv_cache_events and not KV_EVENTS_AVAILABLE: + logger.warning( + "KV cache events are enabled (enable_kv_cache_events=True) " + "but KV Events module is not available. Events will not be emitted." + ) + + self.reset() + + ##### Public API ##### + + def reset(self): + self.root_node = TreeNode() + self.root_node.key = [] + self.root_node.value = [] + self.root_node.host_value = [] + self.root_node.lock_ref = 1 + self.root_node.hash_value = [] + self.evictable_size_ = 0 + self.protected_size_ = 0 + self.evictable_leaves: set[TreeNode] = set() + + if self.enable_kv_cache_events and KV_EVENTS_AVAILABLE: + self.kv_event_queue.append(AllBlocksCleared()) + + def match_prefix(self, key: list, **kwargs) -> MatchResult: + """Find the matching prefix from the prefix tree. + Args: + key: A list of token IDs to find a matching prefix. + Returns: + A tuple of a tensor of matching prefix token IDs and + the last node that contains the prefix values. Note that + this API can modify the internal state of the prefix tree. + The last node creates a new child if the prefix is shorter + than the last node's value. + """ + + if self.disable or len(key) == 0: + return self._empty_match_result() + + page_size = self.token_to_kv_pool_allocator.page_size + # Compatible with whether the incoming key is paged + if not isinstance(key[0], tuple): + full_page_num = len(key) // page_size + paged_token_ids = [ + tuple(key[i * page_size : (i + 1) * page_size]) + for i in range(0, full_page_num) + ] + else: + paged_token_ids = key + if len(paged_token_ids) == 0: + return self._empty_match_result() + value = [] + last_node = [self.root_node] + self._match_prefix_helper(self.root_node, paged_token_ids, value, last_node) + if value and isinstance(value[0], list): + flat_value = [e for arr in value for e in arr] + value = torch.concat(flat_value) + elif value: + value = torch.concat(value) + else: + value = torch.tensor([], dtype=torch.int32) + return MatchResult( + device_indices=value, + last_device_node=last_node[0], + last_host_node=last_node[0], + device_prefix_length=len(value) * page_size, + ) + + def insert(self, key: list, value=None): + if self.disable: + return 0 + + if value is None: + value = [x for x in key] + return self._insert_helper(self.root_node, key, value) + + def cache_finished_req( + self, + req: Req, + token_ids: list[int] | None = None, + delay_req_pool_release=False, + ): + """Cache request when it finishes.""" + if self.disable: + alloced_len = self.req_to_token_pool.alloced_lens[req.req_pool_idx].item() + self.token_to_kv_pool_allocator.free_req_cache( + req.req_pool_idx, alloced_len + ) + if not delay_req_pool_release: + self.req_to_token_pool.free(req.req_pool_idx) + req.req_pool_idx = None + return + + if token_ids is None: + token_ids = (req.origin_input_ids + req.output_ids)[:-1] + # Prefix Cache takes one ref in memory pool + req_pool_idx = req.req_pool_idx + page_size = self.token_to_kv_pool_allocator.page_size + seq_len = len(token_ids) + # without last not full page + full_page_num = seq_len // page_size + paged_token_ids = [ + tuple(token_ids[i * page_size : (i + 1) * page_size]) + for i in range(0, full_page_num) + ] + page_ids = self.token_to_kv_pool_allocator.req_to_page[ + req_pool_idx, :full_page_num + ].clone() + logger.debug("insert_page_ids=%s", page_ids) + _ = self.insert(paged_token_ids, page_ids) + + # After insert, the tree changed so the match prefix ids changed + new_prefix_page_ids = self.match_prefix(paged_token_ids).device_indices + # new_prefix_page_ids is cached in tree, free the diff part in page_ids + + if new_prefix_page_ids.numel() > 0: + self.token_to_kv_pool_allocator.free_with_diff( + new_prefix_page_ids, page_ids + ) + self.token_to_kv_pool_allocator.free_extra_pages_not_cached( + req_pool_idx, + seq_len, + self.req_to_token_pool.alloced_lens[req_pool_idx].item(), + ) + if not delay_req_pool_release: + self.req_to_token_pool.free(req.req_pool_idx) + req.req_pool_idx = None + # Remove req slot release the cache lock + if req.last_node is not None: + self.dec_lock_ref(req.last_node) + req.last_node = None + logger.debug( + "[cache_finished_req]\nold_prefix_page_ids=%s\nnew_prefix_page_ids=%s\nlen(token_ids)=%s \nalloced_lens=%s, self.evictable_size_=%r, self.protected_size_=%r, req.last_node=%r", + self.token_to_kv_pool_allocator.req_to_page[ + req_pool_idx, : (seq_len + page_size - 1) // page_size + ].tolist(), + new_prefix_page_ids.tolist(), + len(token_ids), + self.req_to_token_pool.alloced_lens[req_pool_idx], + self.evictable_size_, + self.protected_size_, + req.last_node, + ) + + def cache_unfinished_req(self, req: Req, token_ids: list[int] | None = None): + """Cache request when it is unfinished.""" + + if self.disable: + return + + if token_ids is None: + token_ids = req.fill_ids + + page_size = self.token_to_kv_pool_allocator.page_size + req_pool_idx = req.req_pool_idx + seq_len = len(token_ids) + # without last not full page + full_page_num = seq_len // page_size + paged_token_ids = [ + tuple(token_ids[i * page_size : (i + 1) * page_size]) + for i in range(0, full_page_num) + ] + + page_ids = self.token_to_kv_pool_allocator.req_to_page[ + req_pool_idx, :full_page_num + ].clone() + _ = self.insert(paged_token_ids, page_ids) + # After insert, perform matching, use page_id in prefix tree to replace the allocated + # page before insert, release diff part, and write to req_to_token_pool + match_result = self.match_prefix(paged_token_ids) + new_prefix_page_ids, new_last_node = ( + match_result.device_indices, + match_result.last_device_node, + ) + if new_prefix_page_ids.numel() > 0: + diff = self.token_to_kv_pool_allocator.free_with_diff( + new_prefix_page_ids, page_ids + ) + diff_idxs = torch.nonzero(diff) + self.token_to_kv_pool_allocator.req_to_page[ + req.req_pool_idx, diff_idxs.squeeze() + ] = new_prefix_page_ids[diff] + # Fix requires prefix cache to store CPU page IDs alongside GPU ones. + self.token_to_kv_pool_allocator.req_to_page_cpu[ + req.req_pool_idx, diff_idxs.squeeze().cpu() + ] = new_prefix_page_ids[diff].cpu() + token_level_offsets = torch.arange( + self.page_size, device=self.req_to_token_pool.device + ) + indices_start_locs = diff_idxs * self.page_size + diff_slots_indices = ( + indices_start_locs[:, None] + token_level_offsets + ).flatten() + new_slots_start_locs = new_prefix_page_ids[diff] * self.page_size + new_slots = (new_slots_start_locs[:, None] + token_level_offsets).flatten() + self.req_to_token_pool.req_to_token[ + req.req_pool_idx, diff_slots_indices + ] = new_slots.to(torch.int32) + self.dec_lock_ref(req.last_node) + self.inc_lock_ref(new_last_node) + req.last_node = new_last_node + + def pretty_print(self, start_str: str = ""): + logger.debug( + self._print_helper( + self.root_node, 0, start_str + f" #tokens: {self.total_size()} " + ) + ) + + def total_size(self): + return self._total_size_helper(self.root_node) + + def evict(self, num_tokens: int, evict_callback: Callable = None): + if self.disable: + return + + heap = [ + (self.eviction_strategy.get_priority(node), node) + for node in self.evictable_leaves + ] + heapq.heapify(heap) + + num_evicted = 0 + while num_evicted < num_tokens and heap: + _, x = heapq.heappop(heap) + + # evictable_leaves only contains unlocked leaves so lock_ref > 0 can + # only happen for cascade parents pushed mid-loop; guard defensively. + if x.lock_ref > 0: + continue + + self.token_to_kv_pool_allocator.append_to_later_free(x.value) + num_evicted += len(x.value) + self._delete_leaf(x) # removes x from evictable_leaves, may add parent + + # Push cascade parent onto the working heap so it can be evicted in + # this same call. _delete_leaf already added it to evictable_leaves. + parent = x.parent + if ( + not parent.children + and parent.lock_ref == 0 + and parent != self.root_node + ): + heapq.heappush( + heap, (self.eviction_strategy.get_priority(parent), parent) + ) + + self.token_to_kv_pool_allocator.free_group_end() + return num_evicted + + def inc_lock_ref(self, node: TreeNode): + if self.disable: + return 0 + + delta = 0 + while node != self.root_node: + if node.lock_ref == 0: + self.evictable_size_ -= len(node.value) + self.protected_size_ += len(node.value) + delta -= len(node.value) + self.evictable_leaves.discard(node) + node.lock_ref += 1 + node = node.parent + return delta + + def dec_lock_ref(self, node: TreeNode): + if self.disable: + return 0 + + if node is None: + return 0 + + delta = 0 + while node != self.root_node: + if node.lock_ref == 1: + self.evictable_size_ += len(node.value) + self.protected_size_ -= len(node.value) + delta += len(node.value) + if not node.children: + self.evictable_leaves.add(node) + node.lock_ref -= 1 + node = node.parent + return delta + + def evictable_size(self): + return self.evictable_size_ + + def protected_size(self): + # protected size refers to the size of the cache that is locked + return self.protected_size_ + + ##### Internal Helper Functions ##### + + def _update_leaf_status(self, node: TreeNode) -> None: + if node == self.root_node: + return + if not node.children and node.lock_ref == 0: + self.evictable_leaves.add(node) + else: + self.evictable_leaves.discard(node) + + def _empty_match_result(self): + return MatchResult( + device_indices=torch.tensor([], dtype=torch.int32), + device_prefix_length=0, + host_hit_length=0, + last_device_node=self.root_node, + last_host_node=self.root_node, + ) + + def _match_prefix_helper( + self, node: TreeNode, key: list, value, last_node: TreeNode + ): + node.last_access_time = time.monotonic() + if not key: + return + + if key[0] in node.children: + child = node.children[key[0]] + prefix_len = _key_match(child.key, key) + if prefix_len < len(child.key): + new_node = self._split_node(child.key, child, prefix_len) + value.append(new_node.value) + last_node[0] = new_node + else: + value.append(child.value) + last_node[0] = child + self._match_prefix_helper(child, key[prefix_len:], value, last_node) + + def _split_node(self, key, child: TreeNode, split_len: int): + # new_node -> child + new_node = TreeNode() + new_node.children = {key[split_len]: child} + new_node.parent = child.parent + new_node.lock_ref = child.lock_ref + new_node.key = child.key[:split_len] + new_node.value = child.value[:split_len] + child.parent = new_node + child.key = child.key[split_len:] + child.value = child.value[split_len:] + new_node.parent.children[key[0]] = new_node + return new_node + + def _insert_helper(self, node: TreeNode, key: list, value): + node.last_access_time = time.monotonic() + if not key: + return 0 + + if key[0] in node.children: + child = node.children[key[0]] + prefix_len = _key_match(child.key, key) + + if prefix_len == len(child.key): + if prefix_len == len(key): + return prefix_len + else: + key = key[prefix_len:] + value = value[prefix_len:] + return prefix_len + self._insert_helper(child, key, value) + + new_node = self._split_node(child.key, child, prefix_len) + return prefix_len + self._insert_helper( + new_node, key[prefix_len:], value[prefix_len:] + ) + + if key: + new_node = TreeNode() + new_node.parent = node + new_node.key = key + new_node.value = value + node.children[key[0]] = new_node + self.evictable_size_ += len(value) + # New node is a leaf; parent may have transitioned from leaf to internal. + self._update_leaf_status(new_node) + self._update_leaf_status(node) + + # Emit BlockStored event when new KV blocks are inserted + if self.enable_kv_cache_events: + self._record_store_event(new_node) + return 0 + + def _print_helper(self, node: TreeNode, indent: int, print_str: str) -> str: + for _, child in node.children.items(): + print_str += "\n" + print_str += " " * indent + print_str += f"{child=} key_len={len(child.key)} " + print_str += f"value={child.value} {child.lock_ref=}" + print_str = self._print_helper( + child, indent=indent + 2, print_str=print_str + ) + return print_str + + def _delete_leaf(self, node): + del node.parent.children[node.key[0]] + self.evictable_size_ -= len(node.key) + self.evictable_leaves.discard(node) + # Parent may have become a childless leaf. + self._update_leaf_status(node.parent) + + # Emit BlockRemoved event when KV blocks are removed + if self.enable_kv_cache_events: + self._record_remove_event(node) + + def _total_size_helper(self, node: TreeNode): + if node.evicted: + return 0 + x = len(node.value) + for child in node.children.values(): + x += self._total_size_helper(child) + return x + + def _collect_leaves(self): + ret_list = [] + stack = [self.root_node] + + while stack: + cur_node = stack.pop() + if not cur_node.children: + ret_list.append(cur_node) + else: + stack.extend(cur_node.children.values()) + + return ret_list + + def _record_store_event(self, node: TreeNode): + """Record BlockStored event for a node.""" + if ( + not self.enable_kv_cache_events + or not KV_EVENTS_AVAILABLE + or node.key is None + ): + return + + # TokenSpeed uses tuples as keys, where each tuple represents a page + # node.key is a list of tuples: [(token1, token2, ...), (token3, token4, ...), ...] + # Each tuple is already a page, so we iterate over them directly + + parent_block_hash = None + if node.parent and node.parent != self.root_node and node.parent.key: + # Use the last page (tuple) of the parent + parent_block_hash = hash(node.parent.key[-1]) + + # Each element in node.key is already a page (tuple of tokens) + for page_tuple in node.key: + if not page_tuple: + continue + + # Convert tuple to list for token_ids + token_ids = list(page_tuple) + block_hash = hash(page_tuple) + + self.kv_event_queue.append( + BlockStored( + block_hashes=[block_hash], + parent_block_hash=parent_block_hash, + token_ids=token_ids, + block_size=len(token_ids), + lora_id=None, + ) + ) + # Chain next chunk to this one + parent_block_hash = block_hash + + def _record_remove_event(self, node: TreeNode): + """Record BlockRemoved event for a node.""" + if ( + not self.enable_kv_cache_events + or not KV_EVENTS_AVAILABLE + or node.key is None + ): + return + + # Create BlockRemoved event for each page + for start in range(0, len(node.key), self.page_size): + page_tokens = node.key[start : start + self.page_size] + if not page_tokens: + continue + block_hash = hash(tuple(page_tokens)) + self.kv_event_queue.append(BlockRemoved(block_hashes=[block_hash])) diff --git a/python/tokenspeed/runtime/cache/req_to_token_pool.py b/python/tokenspeed/runtime/cache/req_to_token_pool.py new file mode 100755 index 0000000..298332d --- /dev/null +++ b/python/tokenspeed/runtime/cache/req_to_token_pool.py @@ -0,0 +1,122 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +""" +Memory pool. + +ReqToTokenPool maps a request to its token locations. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from tokenspeed.runtime.utils import get_colorful_logger +from tokenspeed.runtime.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter + +logger = get_colorful_logger(__name__) + + +@dataclass +class ReqToTokenPoolInfo: + """For chunked prefill""" + + verified_len: int + alloced_len: int + alloced_slots: torch.Tensor + + +class ReqToTokenPool: + """A memory pool that maps a request to its token locations.""" + + def __init__( + self, + size: int, + max_context_len: int, + device: str, + enable_memory_saver: bool, + ): + memory_saver_adapter = TorchMemorySaverAdapter.create( + enable=enable_memory_saver + ) + + self.size = size + self.max_context_len = max_context_len + self.device = device + # Tag as "kv_cache": this per-token page state is invalid once KV is + # discarded, so it is released/restored alongside the KV cache. + with memory_saver_adapter.region(tag="kv_cache", enable_cpu_backup=False): + self.req_to_token = torch.zeros( + (size, max_context_len), dtype=torch.int32, device=device + ) + # verified_lens records the valid historical KV cache length for each request, + # mainly used to determine the KV cache position to write for this computation + self.verified_lens = torch.zeros(size, dtype=torch.int32, device=device) + # alloced_lens records the allocated KV cache length for each request, + # which can be larger than verified_lens, mainly used to determine the KV cache position for this allocation + self.alloced_lens = torch.zeros(size, dtype=torch.int32, device=device) + self.alloced_lens_cpu = torch.zeros(size, dtype=torch.int32, pin_memory=True) + self.free_slots = list(range(size))[1:] + + def set_req_pool_info(self, req_pool_idx: int, metadata: ReqToTokenPoolInfo): + self.verified_lens[req_pool_idx] = metadata.verified_len + self.alloced_lens[req_pool_idx] = metadata.alloced_len + self.alloced_lens_cpu[req_pool_idx] = metadata.alloced_len + self.req_to_token[req_pool_idx, : metadata.alloced_len] = metadata.alloced_slots + + 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] | None: + if need_size > len(self.free_slots): + return None + + select_index = self.free_slots[:need_size] + self.free_slots = self.free_slots[need_size:] + + # During overlap scheduling, after a retracted request frees its req_pool, + # the forward_thread may still modify its verified_lens, causing errors when + # reusing this position. Here we ensure that when req_idx is reused, the corresponding resource is empty. + self.verified_lens[select_index] = 0 + self.alloced_lens[select_index] = 0 + self.alloced_lens_cpu[select_index] = 0 + + return select_index + + def free(self, free_index: int | list[int]) -> None: + free_indices = [free_index] if isinstance(free_index, int) else free_index + self.free_slots.extend(free_indices) + for index in free_indices: + self.verified_lens[index] = 0 + self.alloced_lens[index] = 0 + self.alloced_lens_cpu[index] = 0 + + def clear(self): + # clear method is called during flush_cache + # slot 0 is used as padding in spec_cuda_graph and is not allocated externally + self.free_slots = list(range(self.size))[1:] + self.verified_lens.zero_() + self.alloced_lens.zero_() + self.alloced_lens_cpu.zero_() diff --git a/python/tokenspeed/runtime/cache/storage/__init__.py b/python/tokenspeed/runtime/cache/storage/__init__.py new file mode 100755 index 0000000..b3578e7 --- /dev/null +++ b/python/tokenspeed/runtime/cache/storage/__init__.py @@ -0,0 +1,27 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Storage backend module for TokenSpeed KVStore.""" + +from tokenspeed.runtime.cache.storage.backend_factory import StorageBackendFactory + +__all__ = [ + "StorageBackendFactory", +] diff --git a/python/tokenspeed/runtime/cache/storage/backend_factory.py b/python/tokenspeed/runtime/cache/storage/backend_factory.py new file mode 100755 index 0000000..fe91735 --- /dev/null +++ b/python/tokenspeed/runtime/cache/storage/backend_factory.py @@ -0,0 +1,180 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import importlib +import logging +from typing import Any + +from tokenspeed.runtime.cache.kvstore_storage import ( + KVStoreStorage, + KVStoreStorageConfig, +) + +logger = logging.getLogger(__name__) + + +class StorageBackendFactory: + """Factory for creating storage backend instances with support for dynamic loading.""" + + _registry: dict[str, dict[str, Any]] = {} + + @staticmethod + def _load_backend_class( + module_path: str, class_name: str, backend_name: str + ) -> type[KVStoreStorage]: + """Load and validate a backend class from module path.""" + try: + module = importlib.import_module(module_path) + backend_class = getattr(module, class_name) + if not issubclass(backend_class, KVStoreStorage): + raise TypeError( + f"Backend class {class_name} must inherit from KVStoreStorage" + ) + return backend_class + except ImportError as exc: + raise ImportError( + f"Failed to import backend '{backend_name}' from '{module_path}': {exc}" + ) from exc + except AttributeError as exc: + raise AttributeError( + f"Class '{class_name}' not found in module '{module_path}': {exc}" + ) from exc + + @classmethod + def register_backend(cls, name: str, module_path: str, class_name: str) -> None: + """Register a storage backend with lazy loading. + + Args: + name: Backend identifier + module_path: Python module path containing the backend class + class_name: Name of the backend class + """ + if name in cls._registry: + logger.warning("Backend '%s' is already registered, overwriting", name) + + def loader() -> type[KVStoreStorage]: + """Lazy loader function to import the backend class.""" + return cls._load_backend_class(module_path, class_name, name) + + cls._registry[name] = { + "loader": loader, + "module_path": module_path, + "class_name": class_name, + } + + @classmethod + def create_backend( + cls, + backend_name: str, + storage_config: KVStoreStorageConfig, + mem_pool_host: Any, + **kwargs, + ) -> KVStoreStorage: + """Create a storage backend instance. + Args: + backend_name: Name of the backend to create + storage_config: Storage configuration + mem_pool_host: Memory pool host object + **kwargs: Additional arguments passed to external backends + Returns: + Initialized storage backend instance + Raises: + ValueError: If backend is not registered and cannot be dynamically loaded + ImportError: If backend module cannot be imported + Exception: If backend initialization fails + """ + # First check if backend is already registered + if backend_name in cls._registry: + registry_entry = cls._registry[backend_name] + backend_class = registry_entry["loader"]() + logger.info( + "Creating storage backend '%s' (%s.%s)", + backend_name, + registry_entry["module_path"], + registry_entry["class_name"], + ) + return backend_class(storage_config) + + # Try to dynamically load backend from extra_config + if backend_name == "dynamic" and storage_config.extra_config is not None: + backend_config = storage_config.extra_config + return cls._create_dynamic_backend( + backend_config, storage_config, mem_pool_host, **kwargs + ) + + # Backend not found + raise ValueError( + f"Unknown storage backend '{backend_name}'. " + f"Registered backends: {list(cls._registry)}. " + ) + + @classmethod + def _create_dynamic_backend( + cls, + backend_config: dict[str, Any], + storage_config: KVStoreStorageConfig, + mem_pool_host: Any, + **kwargs, + ) -> KVStoreStorage: + """Create a backend dynamically from configuration.""" + required_fields = ("backend_name", "module_path", "class_name") + missing_fields = [ + field for field in required_fields if field not in backend_config + ] + if missing_fields: + raise ValueError( + "Missing required fields in backend config for 'dynamic' backend: " + f"{missing_fields}" + ) + + backend_name = backend_config["backend_name"] + module_path = backend_config["module_path"] + class_name = backend_config["class_name"] + + try: + # Import the backend class + backend_class = cls._load_backend_class( + module_path, class_name, backend_name + ) + + logger.info( + "Creating dynamic storage backend '%s' (%s.%s)", + backend_name, + module_path, + class_name, + ) + + # Create the backend instance with storage_config + return backend_class(storage_config, kwargs) + except Exception as exc: + logger.error( + "Failed to create dynamic storage backend '%s': %s", backend_name, exc + ) + raise + + +# Register built-in storage backends +StorageBackendFactory.register_backend( + "mooncake", + "tokenspeed.runtime.cache.storage.mooncake_store.mooncake_store", + "MooncakeStore", +) diff --git a/python/tokenspeed/runtime/cache/storage/mooncake_store/mooncake_store.py b/python/tokenspeed/runtime/cache/storage/mooncake_store/mooncake_store.py new file mode 100755 index 0000000..77d2ac6 --- /dev/null +++ b/python/tokenspeed/runtime/cache/storage/mooncake_store/mooncake_store.py @@ -0,0 +1,715 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import json +import logging +import os +import time +import uuid +from dataclasses import dataclass +from typing import Any + +import requests +import torch + +from tokenspeed.runtime.cache.kv_cache_host import HostKVCache +from tokenspeed.runtime.cache.kvstore_storage import ( + KVStoreStorage, + KVStoreStorageConfig, + KVStoreStorageExtraInfo, +) +from tokenspeed.runtime.utils.env import envs + +DEFAULT_LOCAL_BUFFER_SIZE = 16 * 1024 * 1024 # 16 MB +SETUP_TIMEOUT = 600 # 10min + +logger = logging.getLogger(__name__) + + +def _parse_global_segment_size(value: int | str) -> int: + if isinstance(value, int): + return value + if isinstance(value, str): + s = value.strip().lower() + if s.endswith("gb"): + num = s[:-2].strip() + if not num: + raise ValueError( + "Invalid global_segment_size: missing number before 'gb'" + ) + return int(num) * 1024 * 1024 * 1024 + return int(s) + return int(value) + + +@dataclass +class MooncakeStoreConfig: + local_hostname: str + metadata_server: str + global_segment_size: int + protocol: str + device_name: str + master_server_address: str + master_metrics_port: int + check_server: bool + + @staticmethod + def from_file() -> "MooncakeStoreConfig": + """Load the config from a JSON file.""" + if not envs.TOKENSPEED_KVSTORE_MOONCAKE_CONFIG_PATH.is_set(): + raise RuntimeError( + f"Config file path not set. Please set {envs.TOKENSPEED_KVSTORE_MOONCAKE_CONFIG_PATH.name}" + ) + file_path = envs.TOKENSPEED_KVSTORE_MOONCAKE_CONFIG_PATH.value + try: + with open(file_path) as fin: + config = json.load(fin) + except (OSError, json.JSONDecodeError) as exc: + raise RuntimeError( + f"Failed to load config from {file_path}: {exc}" + ) from exc + + if "master_server_address" not in config: + raise ValueError("master_server_address is required in config file") + + return MooncakeStoreConfig( + local_hostname=config.get( + "local_hostname", envs.MOONCAKE_LOCAL_HOSTNAME.default + ), + metadata_server=config.get( + "metadata_server", envs.MOONCAKE_TE_META_DATA_SERVER.default + ), + global_segment_size=_parse_global_segment_size( + config.get( + "global_segment_size", envs.MOONCAKE_GLOBAL_SEGMENT_SIZE.default + ) + ), + protocol=config.get("protocol", envs.MOONCAKE_PROTOCOL.default), + device_name=config.get("device_name", envs.MOONCAKE_DEVICE.default), + master_server_address=config.get("master_server_address"), + master_metrics_port=config.get( + "master_metrics_port", envs.MOONCAKE_MASTER_METRICS_PORT.default + ), + check_server=config.get("check_server", envs.MOONCAKE_CHECK_SERVER.default), + ) + + @staticmethod + def load_from_env() -> "MooncakeStoreConfig": + """Load config from a file specified in the environment variable. + export MOONCAKE_MASTER=10.13.3.232:50051 + export MOONCAKE_PROTOCOL="rdma" + export MOONCAKE_DEVICE="" + export MOONCAKE_TE_META_DATA_SERVER="P2PHANDSHAKE" + """ + # other required environment variables... + if not envs.MOONCAKE_MASTER.is_set(): + raise ValueError("The environment variable 'MOONCAKE_MASTER' is not set.") + + # Prefer the namespaced Mooncake env var, but keep the older + # LOCAL_HOSTNAME fallback working for existing deployments. + if envs.MOONCAKE_LOCAL_HOSTNAME.is_set(): + local_hostname = envs.MOONCAKE_LOCAL_HOSTNAME.value + else: + local_hostname = os.getenv( + "LOCAL_HOSTNAME", envs.MOONCAKE_LOCAL_HOSTNAME.default + ) + + return MooncakeStoreConfig( + local_hostname=local_hostname, + metadata_server=envs.MOONCAKE_TE_META_DATA_SERVER.value, + global_segment_size=_parse_global_segment_size( + envs.MOONCAKE_GLOBAL_SEGMENT_SIZE.value + ), + protocol=envs.MOONCAKE_PROTOCOL.value, + device_name=envs.MOONCAKE_DEVICE.value, + master_server_address=envs.MOONCAKE_MASTER.value, + master_metrics_port=envs.MOONCAKE_MASTER_METRICS_PORT.value, + check_server=envs.MOONCAKE_CHECK_SERVER.value, + ) + + @staticmethod + def load_from_extra_config(extra_config: dict) -> "MooncakeStoreConfig": + """Load config from extra_config dictionary.""" + if "master_server_address" not in extra_config: + raise ValueError("master_server_address is required in extra_config") + + return MooncakeStoreConfig( + local_hostname=extra_config.get( + "local_hostname", envs.MOONCAKE_LOCAL_HOSTNAME.default + ), + metadata_server=extra_config.get( + "metadata_server", envs.MOONCAKE_TE_META_DATA_SERVER.default + ), + global_segment_size=_parse_global_segment_size( + extra_config.get( + "global_segment_size", envs.MOONCAKE_GLOBAL_SEGMENT_SIZE.default + ) + ), + protocol=extra_config.get("protocol", envs.MOONCAKE_PROTOCOL.default), + device_name=extra_config.get("device_name", envs.MOONCAKE_DEVICE.default), + master_server_address=extra_config["master_server_address"], + master_metrics_port=extra_config.get( + "master_metrics_port", envs.MOONCAKE_MASTER_METRICS_PORT.default + ), + check_server=extra_config.get( + "check_server", envs.MOONCAKE_CHECK_SERVER.default + ), + ) + + +class MooncakeStore(KVStoreStorage): + def __init__(self, storage_config: KVStoreStorageConfig = None): + try: + from mooncake.store import MooncakeDistributedStore + except ImportError as exc: + raise ImportError( + "Please install mooncake by following the instructions at " + "https://kvcache-ai.github.io/Mooncake/getting_started/build.html" + "to run TokenSpeed with MooncakeConnector." + ) from exc + + try: + self.store = MooncakeDistributedStore() + + extra_config = ( + getattr(storage_config, "extra_config", None) + if storage_config + else None + ) + # Load configuration with master_server_address prioritized from extra_config if available + if ( + extra_config is not None + and extra_config.get("master_server_address") is not None + ): + # Load from extra_config + self.config = MooncakeStoreConfig.load_from_extra_config(extra_config) + logger.info( + "Mooncake Configuration loaded from extra_config successfully." + ) + elif envs.TOKENSPEED_KVSTORE_MOONCAKE_CONFIG_PATH.is_set(): + # Load from config file + self.config = MooncakeStoreConfig.from_file() + logger.info("Mooncake Configuration loaded from file successfully.") + else: + # Load from environment variables + self.config = MooncakeStoreConfig.load_from_env() + logger.info("Mooncake Configuration loaded from env successfully.") + + tp_scale_factor = 1 if storage_config is None else storage_config.tp_size + + per_tp_global_segment_size = ( + self.config.global_segment_size // tp_scale_factor + ) + + # Check if extra_backend_tag should be passed to MooncakeDistributedStore + self.extra_backend_tag = None + if extra_config and "extra_backend_tag" in extra_config: + self.extra_backend_tag = extra_config["extra_backend_tag"] + logger.info("Using extra_backend_tag: %s", self.extra_backend_tag) + + # Check server status + if self.config.check_server: + self.check_server() + + # Handle JSON device_name configuration + device_name = self.config.device_name + if device_name and device_name.strip().startswith("{"): + try: + device_config = json.loads(device_name) + if storage_config and hasattr(storage_config, "tp_rank"): + tp_rank = storage_config.tp_rank + # Try both integer and string keys since JSON parsing may convert keys + device_name = device_config.get(tp_rank) or device_config.get( + str(tp_rank), "" + ) + else: + device_name = "" + except (json.JSONDecodeError, AttributeError): + logger.warning( + "Failed to parse device_name as JSON: %s", device_name + ) + device_name = "" + + ret_code = self.store.setup( + self.config.local_hostname, + self.config.metadata_server, + per_tp_global_segment_size, + DEFAULT_LOCAL_BUFFER_SIZE, # Zero copy interface does not need local buffer + self.config.protocol, + device_name, + self.config.master_server_address, + ) + if ret_code: + raise RuntimeError( + f"Failed to setup Mooncake store, error code: {ret_code}" + ) + logger.info("Mooncake store setup successfully.") + + self.warmup() + logger.info("Mooncake store warmup successfully.") + + if storage_config is not None: + self.is_mla_backend = storage_config.is_mla_model + self.local_rank = storage_config.tp_rank + else: + self.is_mla_backend = False + self.local_rank = 0 + + except ValueError as exc: + logger.error("Configuration loading failed: %s", exc) + raise + except Exception as exc: + logger.error("An error occurred while loading the configuration: %s", exc) + raise + + def check_server(self): + master_server_ip, _, _ = self.config.master_server_address.partition(":") + segments_url = f"http://{master_server_ip}:{self.config.master_metrics_port}/get_all_segments" + start_time = time.perf_counter() + + check_result = False + while time.perf_counter() - start_time < SETUP_TIMEOUT: + try: + check_segments_resp = requests.get(segments_url, timeout=3) + except requests.RequestException: + logger.info( + "waiting mooncake store server started, cost_time: %.2f seconds.", + time.perf_counter() - start_time, + ) + time.sleep(3) + continue + + if check_segments_resp.text == "": + logger.info( + "waiting mooncake store server started, cost_time: %.2f seconds.", + time.perf_counter() - start_time, + ) + time.sleep(3) + continue + + logger.info("Mooncake store server started successfully.") + check_result = True + break + + if not check_result: + logger.error("Launch mooncake store server timeout") + raise ValueError("Launch mooncake store server timeout") + + def warmup(self): + warmup_key = "tokenspeed_mooncake_store_warmup_key" + uuid.uuid4().hex + warmup_value = bytes(4 * 1024) # 4 KB + put_result = self.store.put(warmup_key, warmup_value) + if put_result != 0: + logger.warning( + "Mooncake store warmup put failed with code %s, skipping warmup (this is expected when global segment size is 0)", + put_result, + ) + return + if self.store.is_exist(warmup_key) != 1: + raise RuntimeError("Mooncake store warmup key was not persisted") + if self.store.get(warmup_key) != warmup_value: + raise RuntimeError("Mooncake store warmup value mismatch") + + def register_mem_pool_host(self, mem_pool_host: HostKVCache): + super().register_mem_pool_host(mem_pool_host) + if self.mem_pool_host.layout not in ("page_first", "page_head"): + raise ValueError( + "mooncake store storage backend only supports page_first or page_head layout" + ) + buffer = self.mem_pool_host.kv_buffer + try: + buffer_ptr = buffer.data_ptr() + buffer_size = buffer.numel() * buffer.element_size() + ret_code = self.store.register_buffer(buffer_ptr, buffer_size) + if ret_code: + logger.error("Failed to register buffer, error code: %s", ret_code) + raise RuntimeError( + f"Failed to register buffer to Mooncake Store, error code: {ret_code}" + ) + except TypeError as err: + logger.error("Failed to register buffer to Mooncake Store: %s", err) + raise TypeError("Mooncake Store Register Buffer Error.") from err + + def _get_mha_buffer_meta(self, keys, indices): + ptr_list, element_size_list = self.mem_pool_host.get_page_buffer_meta(indices) + key_list = [] + for key_ in keys: + key_list.append(f"{key_}_{self.local_rank}_k") + key_list.append(f"{key_}_{self.local_rank}_v") + if len(key_list) != len(ptr_list): + raise ValueError("MHA key metadata does not match buffer metadata") + return key_list, ptr_list, element_size_list + + def _expand_query_keys(self, keys: list[str]) -> tuple[list[str], int, bool]: + if self.is_mla_backend: + return ( + ( + keys + if keys and keys[0].endswith("_k") + else [f"{key}_k" for key in keys] + ), + 1, + False, + ) + + keys_have_suffix = bool( + keys and (keys[0].endswith("_k") or keys[0].endswith("_v")) + ) + if keys_have_suffix: + return keys, 2, False + + query_keys = [ + suffix_key + for key in keys + for suffix_key in ( + f"{key}_{self.local_rank}_k", + f"{key}_{self.local_rank}_v", + ) + ] + return query_keys, 2, True + + @staticmethod + def _expand_pairwise_metadata(values: list[Any], key_count: int) -> list[Any]: + expanded = [] + for index in range(key_count): + expanded.extend((values[index * 2], values[index * 2 + 1])) + return expanded + + def _get_mla_buffer_meta(self, keys, indices): + ptr_list, element_size_list = self.mem_pool_host.get_page_buffer_meta(indices) + key_list = [] + for key_ in keys: + key_list.append(f"{key_}_k") + if len(key_list) != len(ptr_list): + raise ValueError("MLA key metadata does not match buffer metadata") + return key_list, ptr_list, element_size_list + + def _batch_preprocess(self, keys, host_indices): + if not keys: + raise ValueError("keys must not be empty") + if len(keys) != len(host_indices) // self.mem_pool_host.page_size: + raise ValueError("keys length must match host_indices page count") + if self.is_mla_backend: + return self._get_mla_buffer_meta(keys, host_indices) + else: + return self._get_mha_buffer_meta(keys, host_indices) + + def _batch_postprocess(self, results: list[int], is_set_operate=False): + """ + refer to https://github.com/kvcache-ai/Mooncake/blob/main/mooncake-store/include/pybind_client.h + for batch_get_into, results is Vector of integers, + where each element is the number of bytes read on success, or a negative value on error + for batch_put_from, results is Vector of integers, + where each element is 0 on success, or a negative value on error + """ + if self.is_mla_backend: + return [result == 0 if is_set_operate else result > 0 for result in results] + + kv_pairs = zip(results[::2], results[1::2]) + return [ + ( + (k_res == 0 and v_res == 0) + if is_set_operate + else (k_res > 0 and v_res > 0) + ) + for k_res, v_res in kv_pairs + ] + + def batch_get_v1( + self, + keys: list[str], + host_indices: torch.Tensor, + extra_info: KVStoreStorageExtraInfo | None = None, + ) -> list[bool]: + # Apply extra_backend_tag prefix if available + if self.extra_backend_tag is not None: + prefix = self.extra_backend_tag + keys = [f"{prefix}_{key}" for key in keys] + + key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices) + get_results = self._get_batch_zero_copy_impl( + key_strs, buffer_ptrs, buffer_sizes + ) + return self._batch_postprocess(get_results, is_set_operate=False) + + def batch_set_v1( + self, + keys: list[str], + host_indices: torch.Tensor, + extra_info: KVStoreStorageExtraInfo | None = None, + ) -> list[bool]: + # Apply extra_backend_tag prefix if available + if self.extra_backend_tag is not None: + prefix = self.extra_backend_tag + keys = [f"{prefix}_{key}" for key in keys] + + key_strs, buffer_ptrs, buffer_sizes = self._batch_preprocess(keys, host_indices) + exist_result = self._batch_exist(key_strs) + + set_keys = [] + set_buffer_ptrs = [] + set_buffer_sizes = [] + set_indices = [] + set_results = [-1] * len(key_strs) + for index, key_str in enumerate(key_strs): + if exist_result[index] != 1: + set_keys.append(key_str) + set_buffer_ptrs.append(buffer_ptrs[index]) + set_buffer_sizes.append(buffer_sizes[index]) + set_indices.append(index) + else: + set_results[index] = 0 + + # Only set non-existing keys to storage + if set_keys: + put_results = self._put_batch_zero_copy_impl( + set_keys, set_buffer_ptrs, set_buffer_sizes + ) + for index, set_index in enumerate(set_indices): + set_results[set_index] = put_results[index] + + return self._batch_postprocess(set_results, is_set_operate=True) + + def set( + self, + key, + value: Any | None = None, + target_location: list[int] | None = None, + target_sizes: list[int] | None = None, + ) -> bool: + # Only support zero copy set for now + if target_location is None or target_sizes is None: + raise ValueError("target_location and target_sizes are required") + # Format key with local_rank suffix for non-MLA backend + if self.is_mla_backend: + query_keys = [f"{key}_k"] + target_locations = [target_location] + target_sizes_list = [target_sizes] + else: + # For non-MLA backend, we need to set both k and v + query_keys = [f"{key}_{self.local_rank}_k", f"{key}_{self.local_rank}_v"] + # target_location and target_sizes should be lists with 2 elements + if isinstance(target_location, list) and len(target_location) >= 2: + target_locations = [target_location[0], target_location[1]] + else: + # If not a list, assume it's a single location for k only + target_locations = [target_location, target_location] + + if isinstance(target_sizes, list) and len(target_sizes) >= 2: + target_sizes_list = [target_sizes[0], target_sizes[1]] + else: + # If not a list, assume it's a single size for k only + target_sizes_list = [target_sizes, target_sizes] + + exist_result = self._batch_exist(query_keys) + set_keys = [] + set_target_locations = [] + set_target_sizes = [] + for index, query_key in enumerate(query_keys): + if exist_result[index] != 1: + set_keys.append(query_key) + set_target_locations.append(target_locations[index]) + set_target_sizes.append(target_sizes_list[index]) + + # Only set non-existing keys to storage + if set_keys: + put_result = self._put_batch_zero_copy_impl( + set_keys, set_target_locations, set_target_sizes + ) + if any(result != 0 for result in put_result): + return False + return True + + def batch_set( + self, + keys: list[str], + values: list[torch.Tensor] | None = None, + target_locations: list[int] | None = None, + target_sizes: list[int] | None = None, + ) -> bool: + # Only support zero copy set for now + if target_locations is None or target_sizes is None: + raise ValueError("target_locations and target_sizes are required") + if len(keys) != len(target_locations) or len(keys) != len(target_sizes): + raise ValueError( + "keys, target_locations, and target_sizes must have matching lengths" + ) + + if not keys: + return False + + if any( + key is None or location is None or size is None + for key, location, size in zip(keys, target_locations, target_sizes) + ): + return False + + query_keys, _, expanded_non_mla = self._expand_query_keys(keys) + if expanded_non_mla: + expanded_target_locations = self._expand_pairwise_metadata( + target_locations, len(keys) + ) + expanded_target_sizes = self._expand_pairwise_metadata( + target_sizes, len(keys) + ) + else: + expanded_target_locations = target_locations + expanded_target_sizes = target_sizes + + exist_result = self._batch_exist(query_keys) + set_keys = [] + set_target_locations = [] + set_target_sizes = [] + set_indices = [] + for index, query_key in enumerate(query_keys): + if exist_result[index] != 1: + set_keys.append(query_key) + set_target_locations.append(expanded_target_locations[index]) + set_target_sizes.append(expanded_target_sizes[index]) + set_indices.append(index) + # Only set non-existing keys to storage + + put_result = self._put_batch_zero_copy_impl( + set_keys, set_target_locations, set_target_sizes + ) + for index, set_index in enumerate(set_indices): + if put_result[index] == 0: + exist_result[set_index] = 1 + + success_count = 0 + for index in range(len(query_keys)): + if exist_result[index] == 0: + break + success_count += 1 + return success_count == len(query_keys) + + def get( + self, + key, + target_location: Any | None = None, + target_sizes: Any | None = None, + ) -> bool: + if target_location is None or target_sizes is None: + raise ValueError("target_location and target_sizes are required") + # Format key with local_rank suffix for non-MLA backend + if self.is_mla_backend: + query_keys = [f"{key}_k"] + target_locations = [target_location] + target_sizes_list = [target_sizes] + else: + # For non-MLA backend, we need to get both k and v + query_keys = [f"{key}_{self.local_rank}_k", f"{key}_{self.local_rank}_v"] + # target_location and target_sizes should be lists with 2 elements + if isinstance(target_location, list) and len(target_location) >= 2: + target_locations = [target_location[0], target_location[1]] + else: + # If not a list, assume it's a single location for k only + target_locations = [target_location, target_location] + + if isinstance(target_sizes, list) and len(target_sizes) >= 2: + target_sizes_list = [target_sizes[0], target_sizes[1]] + else: + # If not a list, assume it's a single size for k only + target_sizes_list = [target_sizes, target_sizes] + + get_result = self._get_batch_zero_copy_impl( + query_keys, target_locations, target_sizes_list + ) + # Return True only if both k and v are successfully retrieved + return all(result >= 0 for result in get_result) + + def batch_get( + self, + keys: list[str], + target_locations: Any | None = None, + target_sizes: Any | None = None, + ) -> int: + if target_locations is None or target_sizes is None: + raise ValueError("target_locations and target_sizes are required") + if len(keys) != len(target_locations) or len(keys) != len(target_sizes): + raise ValueError( + "keys, target_locations, and target_sizes must have matching lengths" + ) + if not keys: + return 0 + + query_keys, key_multiplier, expanded_non_mla = self._expand_query_keys(keys) + + # Note: target_locations and target_sizes need to match the query_keys length + # If keys already have suffixes, target_locations and target_sizes should already match + # If keys don't have suffixes, we need to expand them for non-MLA backend + if expanded_non_mla: + # Expand target_locations and target_sizes to match query_keys + target_locations = self._expand_pairwise_metadata( + target_locations, len(keys) + ) + target_sizes = self._expand_pairwise_metadata(target_sizes, len(keys)) + + get_result = self._get_batch_zero_copy_impl( + query_keys, target_locations, target_sizes + ) + + for index in range(len(query_keys)): + if get_result[index] < 0: + return index // key_multiplier + return len(query_keys) // key_multiplier + + def exists(self, key) -> bool: + # Format key with local_rank suffix for non-MLA backend + if self.is_mla_backend: + query_keys = [f"{key}_k"] + else: + # For non-MLA backend, we need to check both k and v + query_keys = [f"{key}_{self.local_rank}_k", f"{key}_{self.local_rank}_v"] + exist_result = self._batch_exist(query_keys) + return all(result == 1 for result in exist_result) + + def batch_exists( + self, keys, extra_info: KVStoreStorageExtraInfo | None = None + ) -> int: + query_keys, key_multiplier, _ = self._expand_query_keys(keys) + + exist_result = self._batch_exist(query_keys) + + for index in range(len(query_keys)): + if exist_result[index] != 1: + return index // key_multiplier + return len(query_keys) // key_multiplier + + def close(self): + # MooncakeDistributedStore will automatically call the destructor, so + # it is unnecessary to close it manually. + return None + + def clear(self) -> None: + self.store.remove_all() + + def _put_batch_zero_copy_impl( + self, key_strs: list[str], buffer_ptrs: list[int], buffer_sizes: list[int] + ) -> list[int]: + return self.store.batch_put_from(key_strs, buffer_ptrs, buffer_sizes) + + def _get_batch_zero_copy_impl( + self, key_strs: list[str], buffer_ptrs: list[int], buffer_sizes: list[int] + ) -> list[int]: + return self.store.batch_get_into(key_strs, buffer_ptrs, buffer_sizes) + + def _batch_exist(self, key_strs: list[str]) -> list[int]: + return self.store.batch_is_exist(key_strs) diff --git a/python/tokenspeed/runtime/cache/transfer/__init__.py b/python/tokenspeed/runtime/cache/transfer/__init__.py new file mode 100644 index 0000000..7c97ee8 --- /dev/null +++ b/python/tokenspeed/runtime/cache/transfer/__init__.py @@ -0,0 +1,39 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from tokenspeed.runtime.cache.transfer.kv_pool import KVCachePool +from tokenspeed.runtime.cache.transfer.mamba_pool import MambaCachePool +from tokenspeed.runtime.cache.transfer.pool import CachePool +from tokenspeed.runtime.cache.transfer.types import ( + CacheKind, + Location, + TransferBatch, + TransferUnit, +) + +__all__ = [ + "CacheKind", + "CachePool", + "KVCachePool", + "Location", + "MambaCachePool", + "TransferBatch", + "TransferUnit", +] diff --git a/python/tokenspeed/runtime/cache/transfer/kv_pool.py b/python/tokenspeed/runtime/cache/transfer/kv_pool.py new file mode 100644 index 0000000..d7d951f --- /dev/null +++ b/python/tokenspeed/runtime/cache/transfer/kv_pool.py @@ -0,0 +1,124 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import torch + +from tokenspeed.runtime.cache.kvstore_controller import LayerDoneCounter +from tokenspeed.runtime.cache.transfer.types import CacheKind + + +class KVCachePool: + kind = CacheKind.KV + + def __init__( + self, + device_pool, + host_pool, + io_backend: str, + layer_num: int, + draft_device_pool=None, + draft_host_pool=None, + draft_layer_num: int = 0, + ): + self.device_pool = device_pool + self.host_pool = host_pool + self.io_backend = io_backend + self.layer_num = layer_num + self.draft_device_pool = draft_device_pool + self.draft_host_pool = draft_host_pool + self.draft_layer_num = draft_layer_num + self._counter = LayerDoneCounter(max(layer_num, draft_layer_num, 1)) + device_pool.register_layer_transfer_counter(self._counter) + + @property + def device(self) -> torch.device | str: + return self.device_pool.device + + @property + def host_layout(self) -> str: + return self.host_pool.layout + + def page_size(self) -> int: + return self.host_pool.page_size + + def num_layers(self) -> int: + return max(self.layer_num, self.draft_layer_num) + + def supports_layerwise_loadback(self) -> bool: + return True + + def get_layer_done_counter(self) -> LayerDoneCounter: + return self._counter + + def local_layer_idx(self, global_layer_id: int) -> int: + return global_layer_id + + def writeback( + self, + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + block_quota: int | None = None, + ) -> None: + self.host_pool.backup_from_device_all_layer( + self.device_pool, + dst_indices, + src_indices, + self.io_backend, + block_quota=block_quota, + ) + if self.draft_host_pool is not None: + self.draft_host_pool.backup_from_device_all_layer( + self.draft_device_pool, + dst_indices, + src_indices, + self.io_backend, + block_quota=block_quota, + ) + + def loadback( + self, src_indices: torch.Tensor, dst_indices: torch.Tensor, layer_idx: int + ) -> None: + if layer_idx < self.layer_num: + self.host_pool.load_to_device_per_layer( + self.device_pool, + src_indices, + dst_indices, + layer_idx, + self.io_backend, + ) + if self.draft_host_pool is not None and layer_idx < self.draft_layer_num: + self.draft_host_pool.load_to_device_per_layer( + self.draft_device_pool, + src_indices, + dst_indices, + layer_idx, + self.io_backend, + ) + + def alloc_host(self, n: int) -> torch.Tensor | None: + return self.host_pool.alloc(n) + + def free_host(self, indices: torch.Tensor) -> None: + self.host_pool.free(indices) + + def host_available(self) -> int: + return self.host_pool.available_size() diff --git a/python/tokenspeed/runtime/cache/transfer/mamba_pool.py b/python/tokenspeed/runtime/cache/transfer/mamba_pool.py new file mode 100644 index 0000000..a3145f1 --- /dev/null +++ b/python/tokenspeed/runtime/cache/transfer/mamba_pool.py @@ -0,0 +1,118 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import torch + +from tokenspeed.runtime.cache.kvstore_controller import LayerDoneCounter +from tokenspeed.runtime.cache.mamba_cache_host import MambaPoolHost +from tokenspeed.runtime.cache.transfer.types import CacheKind +from tokenspeed.runtime.layers.attention.backends.hybrid_linear_attn import ( + SimpleMambaPool, +) + + +class MambaCachePool: + kind = CacheKind.MAMBA + + def __init__( + self, + device_pool: SimpleMambaPool, + host_pool: MambaPoolHost, + io_backend: str, + ): + self.device_pool = device_pool + self.host_pool = host_pool + self.io_backend = io_backend + self._counter = LayerDoneCounter(self.num_layers()) + device_pool.register_layer_transfer_counter(self._counter) + + @property + def device(self) -> torch.device | str: + return self.device_pool.device + + @property + def host_layout(self) -> str: + return self.host_pool.layout + + def page_size(self) -> int: + return 1 + + def num_layers(self) -> int: + return int(self.device_pool.conv_state.shape[0]) + + def supports_layerwise_loadback(self) -> bool: + return True + + def get_layer_done_counter(self) -> LayerDoneCounter: + return self._counter + + def local_layer_idx(self, global_layer_id: int) -> int: + return self.device_pool.mamba_map[global_layer_id] + + def writeback( + self, + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + block_quota: int | None = None, + ) -> None: + self.host_pool.backup_from_device_all_layer( + self.device_pool, + host_indices=dst_indices, + device_indices=src_indices, + io_backend=self.io_backend, + block_quota=block_quota, + ) + + def loadback( + self, src_indices: torch.Tensor, dst_indices: torch.Tensor, layer_idx: int + ) -> None: + self.host_pool.load_to_device_per_layer( + self.device_pool, + host_indices=src_indices, + device_indices=dst_indices, + layer_idx=layer_idx, + io_backend=self.io_backend, + ) + + def copy_layer( + self, src_indices: torch.Tensor, dst_indices: torch.Tensor, layer_idx: int + ) -> None: + if src_indices.numel() == 0: + return + src_indices = src_indices.to( + device=self.device, dtype=torch.int64, non_blocking=True + ) + dst_indices = dst_indices.to( + device=self.device, dtype=torch.int64, non_blocking=True + ) + for cache in self.device_pool.mamba_cache: + layer = cache[layer_idx] + layer.index_copy_(0, dst_indices, layer.index_select(0, src_indices)) + + def alloc_host(self, n: int) -> torch.Tensor | None: + return self.host_pool.alloc(n) + + def free_host(self, indices: torch.Tensor) -> None: + self.host_pool.free(indices) + + def host_available(self) -> int: + return self.host_pool.available_size() diff --git a/python/tokenspeed/runtime/cache/transfer/pool.py b/python/tokenspeed/runtime/cache/transfer/pool.py new file mode 100644 index 0000000..1ccf348 --- /dev/null +++ b/python/tokenspeed/runtime/cache/transfer/pool.py @@ -0,0 +1,64 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from typing import Protocol + +import torch + +from tokenspeed.runtime.cache.transfer.types import CacheKind + + +class CachePool(Protocol): + kind: CacheKind + device: torch.device | str + host_layout: str + + def page_size(self) -> int: ... + + def num_layers(self) -> int: ... + + def supports_layerwise_loadback(self) -> bool: ... + + def writeback( + self, + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + block_quota: int | None = None, + ) -> None: ... + + def loadback( + self, src_indices: torch.Tensor, dst_indices: torch.Tensor, layer_idx: int + ) -> None: ... + + def copy_layer( + self, src_indices: torch.Tensor, dst_indices: torch.Tensor, layer_idx: int + ) -> None: ... + + def get_layer_done_counter(self): ... + + def local_layer_idx(self, global_layer_id: int) -> int: ... + + def alloc_host(self, n: int) -> torch.Tensor | None: ... + + def free_host(self, indices: torch.Tensor) -> None: ... + + def host_available(self) -> int: ... diff --git a/python/tokenspeed/runtime/cache/transfer/types.py b/python/tokenspeed/runtime/cache/transfer/types.py new file mode 100644 index 0000000..251090f --- /dev/null +++ b/python/tokenspeed/runtime/cache/transfer/types.py @@ -0,0 +1,60 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum + +import torch + + +class CacheKind(str, Enum): + KV = "kv" + MAMBA = "mamba" + + +class Location(str, Enum): + DEVICE = "device" + HOST = "host" + STORAGE = "storage" + + +@dataclass(slots=True) +class TransferUnit: + kind: CacheKind + src_loc: Location + dst_loc: Location + src_indices: torch.Tensor + dst_indices: torch.Tensor + op_id: int + is_retract: bool = False + layerwise_cow_src_indices: torch.Tensor | None = None + layerwise_cow_dst_indices: torch.Tensor | None = None + + @property + def direction(self) -> tuple[Location, Location]: + return (self.src_loc, self.dst_loc) + + +@dataclass(slots=True) +class TransferBatch: + units: list[TransferUnit] + op_ids: list[int] diff --git a/python/tokenspeed/runtime/cache/utils.py b/python/tokenspeed/runtime/cache/utils.py new file mode 100755 index 0000000..fb2ee15 --- /dev/null +++ b/python/tokenspeed/runtime/cache/utils.py @@ -0,0 +1,367 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Common helper utilities for mem-cache operations.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def set_mla_kv_buffer_kernel( + kv_buffer_ptr, + cache_k_nope_ptr, + cache_k_rope_ptr, + loc_ptr, + buffer_stride: tl.constexpr, + nope_stride: tl.constexpr, + rope_stride: tl.constexpr, + nope_dim: tl.constexpr, + rope_dim: tl.constexpr, + BLOCK: tl.constexpr, + ENABLE_PDL: tl.constexpr, +): + if ENABLE_PDL: + tl.extra.cuda.gdc_wait() + + pid_loc = tl.program_id(0) + pid_blk = tl.program_id(1) + + base = pid_blk * BLOCK + offs = base + tl.arange(0, BLOCK) + total_dim = nope_dim + rope_dim + mask = offs < total_dim + + loc = tl.load(loc_ptr + pid_loc).to(tl.int64) + dst_ptr = kv_buffer_ptr + loc * buffer_stride + offs + + if base + BLOCK <= nope_dim: + src = tl.load( + cache_k_nope_ptr + pid_loc * nope_stride + offs, + mask=mask, + ) + else: + offs_rope = offs - nope_dim + src = tl.load( + cache_k_rope_ptr + pid_loc * rope_stride + offs_rope, + mask=mask, + ) + + tl.store(dst_ptr, src, mask=mask) + + if ENABLE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def set_mla_kv_buffer_per_loc_kernel( + kv_buffer_ptr, + cache_k_nope_ptr, + cache_k_rope_ptr, + loc_ptr, + n_loc, + buffer_stride: tl.constexpr, + nope_stride: tl.constexpr, + rope_stride: tl.constexpr, + nope_dim: tl.constexpr, + rope_dim: tl.constexpr, + BLOCK_LOC: tl.constexpr, + ENABLE_PDL: tl.constexpr, +): + """Each CTA writes BLOCK_LOC locs (the full nope+rope span for each). + Grid is ceil(n_loc / BLOCK_LOC). With BLOCK_LOC > 1 each CTA processes + a [BLOCK_LOC, nope_dim] tile, exposing more parallelism / vectorization + width and better amortizing launch overhead at large n_loc. + Pairs with the block-split set_mla_kv_buffer_kernel above: + set_mla_kv_buffer_triton dispatches between them. + """ + if ENABLE_PDL: + tl.extra.cuda.gdc_wait() + + pid = tl.program_id(0) + loc_indices = pid * BLOCK_LOC + tl.arange(0, BLOCK_LOC) + loc_mask = loc_indices < n_loc + locs = tl.load(loc_ptr + loc_indices, mask=loc_mask, other=0).to(tl.int64) + + # Nope tile: [BLOCK_LOC, nope_dim] + nope_offs = tl.arange(0, nope_dim) + src_nope = tl.load( + cache_k_nope_ptr + loc_indices[:, None] * nope_stride + nope_offs[None, :], + mask=loc_mask[:, None], + ) + tl.store( + kv_buffer_ptr + locs[:, None] * buffer_stride + nope_offs[None, :], + src_nope, + mask=loc_mask[:, None], + ) + + # Rope tile: [BLOCK_LOC, rope_dim] + rope_offs = tl.arange(0, rope_dim) + src_rope = tl.load( + cache_k_rope_ptr + loc_indices[:, None] * rope_stride + rope_offs[None, :], + mask=loc_mask[:, None], + ) + tl.store( + kv_buffer_ptr + locs[:, None] * buffer_stride + nope_dim + rope_offs[None, :], + src_rope, + mask=loc_mask[:, None], + ) + + if ENABLE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +def set_mla_kv_buffer_triton( + kv_buffer: torch.Tensor, + loc: torch.Tensor, + cache_k_nope: torch.Tensor, + cache_k_rope: torch.Tensor, + enable_pdl: bool = False, +): + # Dispatch buckets from experiments on B200 GPUs. + # n_loc < 512 : block-split kernel — more CTAs/loc fills SMs at decode + # batch sizes. + # n_loc >= 512 : per-loc kernel — fat tiles saturate bandwidth at + # prefill chunk sizes; (BLOCK_LOC, num_warps, num_stages) + # widens with n_loc. Above 16K each loc has enough + # elements to vectorize at 32 threads (16-byte loads). + n_loc = loc.numel() + nope_dim = cache_k_nope.size(-1) + rope_dim = cache_k_rope.size(-1) + + extra_kwargs = {"launch_pdl": True} if enable_pdl else {} + if n_loc >= 512: + if n_loc >= 16384: + block_loc, num_warps, num_stages = 4, 1, 2 + elif n_loc >= 2048: + block_loc, num_warps, num_stages = 4, 4, 2 + else: + block_loc, num_warps, num_stages = 2, 4, 2 + grid = (triton.cdiv(n_loc, block_loc),) + set_mla_kv_buffer_per_loc_kernel[grid]( + kv_buffer, + cache_k_nope, + cache_k_rope, + loc, + n_loc, + kv_buffer.stride(0), + cache_k_nope.stride(0), + cache_k_rope.stride(0), + nope_dim, + rope_dim, + BLOCK_LOC=block_loc, + ENABLE_PDL=enable_pdl, + num_warps=num_warps, + num_stages=num_stages, + **extra_kwargs, + ) + else: + BLOCK = 256 + if nope_dim % BLOCK != 0: + raise ValueError( + f"nope_dim ({nope_dim}) must be a multiple of BLOCK ({BLOCK})" + ) + grid = (n_loc, triton.cdiv(nope_dim + rope_dim, BLOCK)) + set_mla_kv_buffer_kernel[grid]( + kv_buffer, + cache_k_nope, + cache_k_rope, + loc, + kv_buffer.stride(0), + cache_k_nope.stride(0), + cache_k_rope.stride(0), + nope_dim, + rope_dim, + BLOCK=BLOCK, + ENABLE_PDL=enable_pdl, + **extra_kwargs, + ) + + +@triton.jit +def get_mla_kv_buffer_kernel( + kv_buffer_ptr, + cache_k_nope_ptr, + cache_k_rope_ptr, + loc_ptr, + buffer_stride: tl.constexpr, + nope_stride: tl.constexpr, + rope_stride: tl.constexpr, + nope_dim: tl.constexpr, + rope_dim: tl.constexpr, + BLOCK: tl.constexpr, + ENABLE_PDL: tl.constexpr, +): + """Block-split variant: grid (n_loc, ceil(total_dim/BLOCK)), each CTA reads + BLOCK elements of one source (nope OR rope, never straddling). More CTAs/loc + fills SMs better at small n_loc — mirrors the block-split + set_mla_kv_buffer_kernel. Pairs with get_mla_kv_buffer_per_loc_kernel below: + get_mla_kv_buffer_triton dispatches between them. + + Requires BLOCK to divide nope_dim so each block is purely nope or purely + rope (with masking on the trailing rope block). Wrapper picks BLOCK=128. + """ + if ENABLE_PDL: + tl.extra.cuda.gdc_wait() + + pid_loc = tl.program_id(0) + pid_blk = tl.program_id(1) + + base = pid_blk * BLOCK + offs = base + tl.arange(0, BLOCK) + total_dim = nope_dim + rope_dim + mask = offs < total_dim + + loc = tl.load(loc_ptr + pid_loc).to(tl.int64) + src = tl.load(kv_buffer_ptr + loc * buffer_stride + offs, mask=mask) + + if base + BLOCK <= nope_dim: + tl.store(cache_k_nope_ptr + pid_loc * nope_stride + offs, src, mask=mask) + else: + offs_rope = offs - nope_dim + tl.store(cache_k_rope_ptr + pid_loc * rope_stride + offs_rope, src, mask=mask) + + if ENABLE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +@triton.jit +def get_mla_kv_buffer_per_loc_kernel( + kv_buffer_ptr, + cache_k_nope_ptr, + cache_k_rope_ptr, + loc_ptr, + n_loc, + buffer_stride: tl.constexpr, + nope_stride: tl.constexpr, + rope_stride: tl.constexpr, + nope_dim: tl.constexpr, + rope_dim: tl.constexpr, + BLOCK_LOC: tl.constexpr, + ENABLE_PDL: tl.constexpr, +): + """Each CTA reads BLOCK_LOC locs from kv_buffer (gather) and writes them + contiguously to cache_k_nope / cache_k_rope. Grid is ceil(n_loc / BLOCK_LOC). + Mirror of set_mla_kv_buffer_per_loc_kernel with read/write directions + flipped. get_mla_kv_buffer_triton dispatches between this kernel and the + block-split get_mla_kv_buffer_kernel above based on n_loc. + """ + if ENABLE_PDL: + tl.extra.cuda.gdc_wait() + + pid = tl.program_id(0) + loc_indices = pid * BLOCK_LOC + tl.arange(0, BLOCK_LOC) + loc_mask = loc_indices < n_loc + locs = tl.load(loc_ptr + loc_indices, mask=loc_mask, other=0).to(tl.int64) + + # Nope tile: [BLOCK_LOC, nope_dim] — gather from kv_buffer at locs. + nope_offs = tl.arange(0, nope_dim) + src_nope = tl.load( + kv_buffer_ptr + locs[:, None] * buffer_stride + nope_offs[None, :], + mask=loc_mask[:, None], + ) + tl.store( + cache_k_nope_ptr + loc_indices[:, None] * nope_stride + nope_offs[None, :], + src_nope, + mask=loc_mask[:, None], + ) + + # Rope tile: [BLOCK_LOC, rope_dim] + rope_offs = tl.arange(0, rope_dim) + src_rope = tl.load( + kv_buffer_ptr + locs[:, None] * buffer_stride + nope_dim + rope_offs[None, :], + mask=loc_mask[:, None], + ) + tl.store( + cache_k_rope_ptr + loc_indices[:, None] * rope_stride + rope_offs[None, :], + src_rope, + mask=loc_mask[:, None], + ) + + if ENABLE_PDL: + tl.extra.cuda.gdc_launch_dependents() + + +def get_mla_kv_buffer_triton( + kv_buffer: torch.Tensor, + loc: torch.Tensor, + cache_k_nope: torch.Tensor, + cache_k_rope: torch.Tensor, + enable_pdl: bool = False, +): + # Dispatch buckets from experiments on B200 GPUs. + # n_loc < 512 : block-split kernel — more CTAs/loc fills SMs at decode + # batch sizes. + # n_loc >= 512 : per-loc kernel — fat tiles saturate bandwidth. + # The W=4→W=1 transition lands earlier than for set + # (gather reads benefit from fewer threads / wider + # per-thread elements / extra pipeline stages). + n_loc = loc.numel() + nope_dim = cache_k_nope.size(-1) + rope_dim = cache_k_rope.size(-1) + + extra_kwargs = {"launch_pdl": True} if enable_pdl else {} + if n_loc >= 512: + if n_loc >= 16384: + block_loc, num_warps, num_stages = 8, 1, 2 + elif n_loc >= 2048: + block_loc, num_warps, num_stages = 8, 1, 3 + else: + block_loc, num_warps, num_stages = 2, 4, 2 + grid = (triton.cdiv(n_loc, block_loc),) + get_mla_kv_buffer_per_loc_kernel[grid]( + kv_buffer, + cache_k_nope, + cache_k_rope, + loc, + n_loc, + kv_buffer.stride(0), + cache_k_nope.stride(0), + cache_k_rope.stride(0), + nope_dim, + rope_dim, + BLOCK_LOC=block_loc, + ENABLE_PDL=enable_pdl, + num_warps=num_warps, + num_stages=num_stages, + **extra_kwargs, + ) + else: + BLOCK = 256 + if nope_dim % BLOCK != 0: + raise ValueError( + f"nope_dim ({nope_dim}) must be a multiple of BLOCK ({BLOCK})" + ) + grid = (n_loc, triton.cdiv(nope_dim + rope_dim, BLOCK)) + get_mla_kv_buffer_kernel[grid]( + kv_buffer, + cache_k_nope, + cache_k_rope, + loc, + kv_buffer.stride(0), + cache_k_nope.stride(0), + cache_k_rope.stride(0), + nope_dim, + rope_dim, + BLOCK=BLOCK, + ENABLE_PDL=enable_pdl, + **extra_kwargs, + ) diff --git a/python/tokenspeed/runtime/configs/__init__.py b/python/tokenspeed/runtime/configs/__init__.py new file mode 100755 index 0000000..4640c04 --- /dev/null +++ b/python/tokenspeed/runtime/configs/__init__.py @@ -0,0 +1,50 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Runtime configuration exports.""" + +from tokenspeed.runtime.configs.deepseek_v4_config import DeepseekV4Config +from tokenspeed.runtime.configs.kimi_k2_config import KimiK2Config +from tokenspeed.runtime.configs.kimi_k25_config import KimiK25Config +from tokenspeed.runtime.configs.minimax_m2_config import MiniMaxM2Config +from tokenspeed.runtime.configs.qwen2_config import Qwen2Config +from tokenspeed.runtime.configs.qwen3_5_config import Qwen3_5Config, Qwen3_5MoeConfig +from tokenspeed.runtime.configs.qwen3_asr_config import ( + Qwen3ASRAudioEncoderConfig, + Qwen3ASRConfig, + Qwen3ASRThinkerConfig, +) +from tokenspeed.runtime.configs.qwen3_config import Qwen3Config +from tokenspeed.runtime.configs.qwen3_moe_config import Qwen3MoeConfig + +__all__ = [ + "DeepseekV4Config", + "Qwen2Config", + "Qwen3Config", + "Qwen3MoeConfig", + "Qwen3_5Config", + "Qwen3_5MoeConfig", + "Qwen3ASRAudioEncoderConfig", + "Qwen3ASRConfig", + "Qwen3ASRThinkerConfig", + "MiniMaxM2Config", + "KimiK2Config", + "KimiK25Config", +] diff --git a/python/tokenspeed/runtime/configs/deepseek_v4_cache_spec.py b/python/tokenspeed/runtime/configs/deepseek_v4_cache_spec.py new file mode 100644 index 0000000..efddf43 --- /dev/null +++ b/python/tokenspeed/runtime/configs/deepseek_v4_cache_spec.py @@ -0,0 +1,279 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +from tokenspeed.runtime.configs.paged_cache_spec import PagedCacheGroupSpec + +V4_KERNEL_BLOCK_ROWS: int = 64 +V4_SWA_KV_GROUP_ID = "v4.swa_kv" +V4_INDEXER_COMPRESSOR_STATE_GROUP_ID = "v4.c4a.indexer_compressor_state" +DEEPSEEK_V4_FP8_MAX = 448.0 +DEEPSEEK_V4_FP8_BLOCK_SIZE = 128 +DEEPSEEK_V4_FP8_QUANT_BLOCK = 64 +DEEPSEEK_V4_FP8_INDEXER_BLOCK_SIZE = 128 +DEEPSEEK_V4_FP8_SCALE_BYTES = 4 +DEEPSEEK_V4_MXFP4_BLOCK_SIZE = 32 +DEEPSEEK_V4_MXFP4_SCALE_BYTES = 1 +DEEPSEEK_V4_SPARSE_PREFILL_TOPK_ALIGNMENT = 128 +DEEPSEEK_V4_COMPRESSED_LOGICAL_BLOCK_SIZE = 256 +_COMPRESSOR_STATE_WINDOW_TOKENS = {4: 8, 128: 128} +_COMPRESSOR_STATE_ROWS_PER_PAGE = {4: 4, 128: 8} + + +def deepseek_v4_nope_dim(head_dim: int, rope_dim: int) -> int: + nope_dim = int(head_dim) - int(rope_dim) + if nope_dim <= 0: + raise ValueError(f"head_dim={head_dim} must be larger than rope_dim={rope_dim}") + return nope_dim + + +def deepseek_v4_swa_token_stride(head_dim: int, rope_dim: int) -> int: + return deepseek_v4_nope_dim(head_dim, rope_dim) + int(rope_dim) * 2 + + +def deepseek_v4_swa_scale_dim(head_dim: int, rope_dim: int) -> int: + nope_dim = deepseek_v4_nope_dim(head_dim, rope_dim) + if nope_dim % DEEPSEEK_V4_FP8_QUANT_BLOCK != 0: + raise ValueError( + "DeepSeek V4 FP8 NoPE dim must be divisible by " + f"{DEEPSEEK_V4_FP8_QUANT_BLOCK}, got {nope_dim}" + ) + return nope_dim // DEEPSEEK_V4_FP8_QUANT_BLOCK + 1 + + +def deepseek_v4_swa_row_bytes(head_dim: int, rope_dim: int) -> int: + return deepseek_v4_swa_token_stride(head_dim, rope_dim) + deepseek_v4_swa_scale_dim( + head_dim, rope_dim + ) + + +def deepseek_v4_indexer_mxfp4_value_bytes(index_head_dim: int) -> int: + index_head_dim = int(index_head_dim) + if index_head_dim % 2 != 0: + raise ValueError(f"MXFP4 index head dim must be even, got {index_head_dim}") + return index_head_dim // 2 + + +def deepseek_v4_indexer_mxfp4_scale_dim(index_head_dim: int) -> int: + index_head_dim = int(index_head_dim) + if index_head_dim % DEEPSEEK_V4_MXFP4_BLOCK_SIZE != 0: + raise ValueError( + "MXFP4 index head dim must be divisible by " + f"{DEEPSEEK_V4_MXFP4_BLOCK_SIZE}, got {index_head_dim}" + ) + return ( + index_head_dim // DEEPSEEK_V4_MXFP4_BLOCK_SIZE * DEEPSEEK_V4_MXFP4_SCALE_BYTES + ) + + +def deepseek_v4_indexer_mxfp4_row_bytes(index_head_dim: int) -> int: + return deepseek_v4_indexer_mxfp4_value_bytes( + index_head_dim + ) + deepseek_v4_indexer_mxfp4_scale_dim(index_head_dim) + + +def deepseek_v4_indexer_mxfp4_layout_from_row_bytes( + row_bytes: int, +) -> tuple[int, int, int]: + row_bytes = int(row_bytes) + value_bytes_per_block = DEEPSEEK_V4_MXFP4_BLOCK_SIZE // 2 + bytes_per_block = value_bytes_per_block + DEEPSEEK_V4_MXFP4_SCALE_BYTES + if row_bytes % bytes_per_block != 0: + raise ValueError( + f"MXFP4 indexer row bytes must be value+scale aligned, got {row_bytes}" + ) + num_blocks = row_bytes // bytes_per_block + value_bytes = num_blocks * value_bytes_per_block + scale_bytes = num_blocks * DEEPSEEK_V4_MXFP4_SCALE_BYTES + index_head_dim = num_blocks * DEEPSEEK_V4_MXFP4_BLOCK_SIZE + if deepseek_v4_indexer_mxfp4_scale_dim(index_head_dim) != scale_bytes: + raise ValueError( + f"invalid MXFP4 indexer row bytes {row_bytes} for " + f"index_head_dim={index_head_dim}" + ) + return index_head_dim, value_bytes, scale_bytes + + +def deepseek_v4_indexer_fp8_scale_bytes(index_head_dim: int) -> int: + index_head_dim = int(index_head_dim) + if index_head_dim % DEEPSEEK_V4_FP8_INDEXER_BLOCK_SIZE != 0: + raise ValueError( + "FP8 index head dim must be divisible by " + f"{DEEPSEEK_V4_FP8_INDEXER_BLOCK_SIZE}, got {index_head_dim}" + ) + return ( + index_head_dim + // DEEPSEEK_V4_FP8_INDEXER_BLOCK_SIZE + * DEEPSEEK_V4_FP8_SCALE_BYTES + ) + + +def deepseek_v4_indexer_fp8_row_bytes(index_head_dim: int) -> int: + return int(index_head_dim) + deepseek_v4_indexer_fp8_scale_bytes(index_head_dim) + + +def deepseek_v4_indexer_fp8_layout_from_row_bytes( + row_bytes: int, +) -> tuple[int, int]: + row_bytes = int(row_bytes) + bytes_per_block = DEEPSEEK_V4_FP8_INDEXER_BLOCK_SIZE + DEEPSEEK_V4_FP8_SCALE_BYTES + if row_bytes % bytes_per_block != 0: + raise ValueError( + f"FP8 indexer row bytes must be value+scale aligned, got {row_bytes}" + ) + index_head_dim = row_bytes // bytes_per_block * DEEPSEEK_V4_FP8_INDEXER_BLOCK_SIZE + scale_bytes = deepseek_v4_indexer_fp8_scale_bytes(index_head_dim) + if index_head_dim + scale_bytes != row_bytes: + raise ValueError( + f"invalid FP8 indexer row bytes {row_bytes} for " + f"index_head_dim={index_head_dim}" + ) + return index_head_dim, scale_bytes + + +def v4_compressor_state_group_id(ratio: int) -> str: + return f"v4.c{int(ratio)}a.compressor_state" + + +def v4_compressed_kv_group_id(ratio: int) -> str: + return f"v4.c{int(ratio)}a.compressed_kv" + + +def parse_v4_compressor_state_group_id(group_id: str) -> int | None: + prefix = "v4.c" + suffix = "a.compressor_state" + if not group_id.startswith(prefix) or not group_id.endswith(suffix): + return None + ratio_text = group_id[len(prefix) : -len(suffix)] + try: + return int(ratio_text) + except ValueError: + return None + + +def _compressed_kernel_block_size(ratio: int) -> int: + if ratio <= 1: + raise ValueError(f"ratio must be > 1, got {ratio}") + return max(1, DEEPSEEK_V4_COMPRESSED_LOGICAL_BLOCK_SIZE // ratio) + + +def _resolve_sliding_window(hf_config: Any) -> int: + for source in (hf_config, getattr(hf_config, "text_config", None)): + if source is None: + continue + if hasattr(source, "sliding_window"): + value = source.sliding_window + if value is None: + raise ValueError("DeepSeek V4 sliding_window is None") + window = int(value) + if window <= 0: + raise ValueError(f"sliding_window must be positive, got {value!r}") + return window + raise ValueError("DeepSeek V4 hf_config is missing sliding_window") + + +def build_v4_cache_specs( + hf_config: Any, + *, + layer_ratio: Sequence[int], +) -> list[PagedCacheGroupSpec]: + swa_window = _resolve_sliding_window(hf_config) + unique_compress_ratios = sorted({int(r) for r in layer_ratio if int(r) > 1}) + + specs: list[PagedCacheGroupSpec] = [ + # SWA kv: trailing window only -> State family. + PagedCacheGroupSpec( + group_id=V4_SWA_KV_GROUP_ID, + retention="sliding_window", + rows_per_page=V4_KERNEL_BLOCK_ROWS, + entry_stride_tokens=1, + sliding_window_tokens=swa_window, + family="state", + ), + ] + for ratio in unique_compress_ratios: + if ratio not in _COMPRESSOR_STATE_WINDOW_TOKENS: + raise ValueError(f"unsupported DeepSeek V4 compress_ratio={ratio}") + # Compressor state: tail buffer -> State family. + specs.append( + PagedCacheGroupSpec( + group_id=v4_compressor_state_group_id(ratio), + retention="sliding_window", + rows_per_page=_COMPRESSOR_STATE_ROWS_PER_PAGE[ratio], + entry_stride_tokens=1, + sliding_window_tokens=_COMPRESSOR_STATE_WINDOW_TOKENS[ratio], + family="state", + ) + ) + # Compressed kv: full-history chain (indexer K shares this group). + specs.append( + PagedCacheGroupSpec( + group_id=v4_compressed_kv_group_id(ratio), + retention="full_history", + rows_per_page=_compressed_kernel_block_size(ratio), + entry_stride_tokens=ratio, + sliding_window_tokens=None, + family="history", + ) + ) + if 4 in unique_compress_ratios: + # Indexer compressor state: tail buffer -> State family. + specs.append( + PagedCacheGroupSpec( + group_id=V4_INDEXER_COMPRESSOR_STATE_GROUP_ID, + retention="sliding_window", + rows_per_page=_COMPRESSOR_STATE_ROWS_PER_PAGE[4], + entry_stride_tokens=1, + sliding_window_tokens=_COMPRESSOR_STATE_WINDOW_TOKENS[4], + family="state", + ) + ) + return specs + + +__all__ = [ + "DEEPSEEK_V4_FP8_BLOCK_SIZE", + "DEEPSEEK_V4_COMPRESSED_LOGICAL_BLOCK_SIZE", + "DEEPSEEK_V4_FP8_MAX", + "DEEPSEEK_V4_FP8_INDEXER_BLOCK_SIZE", + "DEEPSEEK_V4_FP8_QUANT_BLOCK", + "DEEPSEEK_V4_FP8_SCALE_BYTES", + "DEEPSEEK_V4_MXFP4_BLOCK_SIZE", + "DEEPSEEK_V4_MXFP4_SCALE_BYTES", + "DEEPSEEK_V4_SPARSE_PREFILL_TOPK_ALIGNMENT", + "V4_INDEXER_COMPRESSOR_STATE_GROUP_ID", + "V4_KERNEL_BLOCK_ROWS", + "V4_SWA_KV_GROUP_ID", + "build_v4_cache_specs", + "deepseek_v4_indexer_fp8_layout_from_row_bytes", + "deepseek_v4_indexer_fp8_row_bytes", + "deepseek_v4_indexer_fp8_scale_bytes", + "deepseek_v4_indexer_mxfp4_layout_from_row_bytes", + "deepseek_v4_indexer_mxfp4_row_bytes", + "deepseek_v4_indexer_mxfp4_scale_dim", + "deepseek_v4_indexer_mxfp4_value_bytes", + "deepseek_v4_nope_dim", + "deepseek_v4_swa_row_bytes", + "deepseek_v4_swa_scale_dim", + "deepseek_v4_swa_token_stride", + "parse_v4_compressor_state_group_id", + "v4_compressed_kv_group_id", + "v4_compressor_state_group_id", +] diff --git a/python/tokenspeed/runtime/configs/deepseek_v4_config.py b/python/tokenspeed/runtime/configs/deepseek_v4_config.py new file mode 100644 index 0000000..3022e35 --- /dev/null +++ b/python/tokenspeed/runtime/configs/deepseek_v4_config.py @@ -0,0 +1,37 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + + +from transformers.configuration_utils import PretrainedConfig + + +class DeepseekV4Config(PretrainedConfig): + model_type = "deepseek_v4" + + def __init__( + self, + max_position_embeddings: int = 1048576, + rope_scaling: dict | None = None, + **kwargs, + ): + self.max_position_embeddings = max_position_embeddings + self.rope_scaling = rope_scaling + self.rope_parameters = rope_scaling or {} + super().__init__(rope_scaling=rope_scaling, **kwargs) diff --git a/python/tokenspeed/runtime/configs/device_config.py b/python/tokenspeed/runtime/configs/device_config.py new file mode 100755 index 0000000..25726b8 --- /dev/null +++ b/python/tokenspeed/runtime/configs/device_config.py @@ -0,0 +1,38 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Runtime device configuration helpers.""" + +import torch + +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +class DeviceConfig: + device: torch.device | None + + def __init__(self, device: str = "cuda") -> None: + if device == "cuda": + self.device_type = device + else: + raise RuntimeError(f"Not supported device type: {device}") + self.device = torch.device(self.device_type) diff --git a/python/tokenspeed/runtime/configs/flat_memory_plan.py b/python/tokenspeed/runtime/configs/flat_memory_plan.py new file mode 100644 index 0000000..0ac38f1 --- /dev/null +++ b/python/tokenspeed/runtime/configs/flat_memory_plan.py @@ -0,0 +1,255 @@ +"""Flat KV-cache memory plan: pure sizing/binding decisions, no torch. + +Components declare per-block bytes as a function of P (block_size): +linear components scale (bytes_per_slot > 0), constant components do not +(const_bytes > 0, mamba state snapshots). Same-(group, layer) components +pack into one page row ([conv|ssm|pad], the vLLM hybrid layout). One +equalizer move: constant rows inflate P until the widest linear row +covers them (vLLM align). plan_tensors then pairs physical slot j with +the j-th layer of every group over a single page-id space and sizes each +slab by its own packed row from the budget. +""" + +from __future__ import annotations + +import math +from collections import defaultdict +from dataclasses import dataclass, replace + +# Labels whose group is state-family (recurrent state rows, not KV history). +# Deliberate one-line duplicate of paged_cache_spec.STATE_LAYER_TYPES: both +# modules are direct-loaded standalone by their tests (importlib, no package +# context), so a cross-module import would break either loader. Keep in sync. +STATE_LAYER_TYPES = frozenset({"linear_attention"}) + + +@dataclass(frozen=True) +class ComponentSpec: + group_id: str + layer: int + component: str + bytes_per_slot: int # linear in P; 0 for constant components + const_bytes: int # constant in P; 0 for linear components + + +@dataclass(frozen=True) +class BlockGeometry: + block_size: int + block_bytes: int + num_blocks: int = 0 # filled by the planners from the memory budget + + +def occurrence_index(labels): + """Within-label occurrence index per position. + + Args: + labels: Iterable of hashable labels (e.g. per-layer type strings). + + Returns: + list[int]: out[i] == number of earlier positions carrying the same + label as position i — the slab pairing order shared by + components_from_layers and the KV pool's slab layout. + """ + counts: dict = {} + out: list[int] = [] + for label in labels: + idx = counts.get(label, 0) + counts[label] = idx + 1 + out.append(idx) + return out + + +def state_const_bytes(conv_shape, conv_dtype, ssm_shape, ssm_dtype): + """Constant per-page state row bytes of one GDN/mamba2 state layer. + + Args: + conv_shape / ssm_shape: Per-layer state tensor shapes (the configs' + mamba2_cache_params conv and temporal shapes). + conv_dtype / ssm_dtype: Matching dtypes (anything with ``itemsize``). + + Returns: + dict[str, int]: {"conv": bytes, "ssm": bytes} — the exact + ``state_const_bytes`` mapping components_from_layers / + equalized_block_size consume (insertion order = row_offset order). + """ + return { + "conv": math.prod(conv_shape) * conv_dtype.itemsize, + "ssm": math.prod(ssm_shape) * ssm_dtype.itemsize, + } + + +def components_from_layers(*, layer_types, kv_bytes_per_slot, state_const_bytes): + """Per-layer ComponentSpecs: history layers carry one linear kv component; + state layers one constant component per state tensor. Layer index is the + within-group occurrence count (the slab pairing order). State component + order (hence row_offset order downstream) follows state_const_bytes + insertion order.""" + comps: list[ComponentSpec] = [] + for label, idx in zip(layer_types, occurrence_index(layer_types)): + if label in STATE_LAYER_TYPES: + for name, nbytes in state_const_bytes.items(): + comps.append(ComponentSpec(label, idx, name, 0, nbytes)) + else: + comps.append(ComponentSpec(label, idx, "kv", kv_bytes_per_slot, 0)) + return comps + + +def _row_demands(components): + """Per-(group, layer) row: (linear bytes-per-slot sum, constant bytes sum).""" + rows = defaultdict(lambda: [0, 0]) + for c in components: + row = rows[(c.group_id, c.layer)] + row[0] += c.bytes_per_slot + row[1] += c.const_bytes + return rows + + +def solve_page_geometry(components, *, block_size, alignment): + """Smallest P >= block_size (multiple of `alignment` when inflated) + such that the widest linear row covers the widest constant row.""" + rows = _row_demands(components).values() + # NOTE: a row mixing linear and constant components is not needed by any + # known model; reject it so the math stays honest. + for lin, const in rows: + if lin > 0 and const > 0: + raise ValueError("a row must be all-linear or all-constant") + max_linear = max((lin for lin, _ in rows), default=0) + max_const = max((const for _, const in rows), default=0) + if max_const > 0: + if max_linear == 0: + raise ValueError("constant components need a linear row to size P against") + needed = -(-max_const // max_linear) # exact integer ceil + if needed > block_size: + block_size = alignment * math.ceil(needed / alignment) + block_bytes = max(max_linear * block_size, max_const) + return BlockGeometry(block_size=block_size, block_bytes=block_bytes) + + +def equalized_block_size( + *, + layer_types, + kv_bytes_per_slot, + state_const_bytes, + block_size, + alignment=None, +): + """Effective P for a state-hybrid profile: `block_size` when the + widest KV row already covers the widest constant state row, else the + smallest multiple of `alignment` that does. `alignment` defaults to the + original `block_size` (the attention backend's page granularity — + no backend declares a finer one), so the inflated P stays a multiple of + the configured block size. Pure wrapper over components_from_layers + + solve_page_geometry so the config-level equalization decision and its + tests share one implementation.""" + comps = components_from_layers( + layer_types=layer_types, + kv_bytes_per_slot=kv_bytes_per_slot, + state_const_bytes=state_const_bytes, + ) + geo = solve_page_geometry( + comps, + block_size=block_size, + alignment=alignment if alignment is not None else block_size, + ) + return geo.block_size + + +@dataclass(frozen=True) +class LayerBinding: + slot: int + group_id: str + layer: int + component: str + nbytes_per_block: int + row_offset: int # byte offset of this component within its (group, layer) page row + + +@dataclass(frozen=True) +class TensorPlan: + name: str + nbytes: int + bindings: tuple[LayerBinding, ...] + + +@dataclass(frozen=True) +class FlatMemoryPlan: + geometry: BlockGeometry + tensors: tuple[TensorPlan, ...] + + +def plan_component_tensors( + components, *, block_size, budget_bytes, reserved_bytes_per_block=0 +): + """One tensor per ComponentSpec, honestly sized: row bytes = that + component's per-block bytes, num_blocks = budget // (sum of all rows + + reserved_bytes_per_block). No cross-component packing, no padding — + every tensor keeps today's standalone-slab shape, so kernels, CUDA + graphs and the host mirror stay untouched. reserved_bytes_per_block + carries co-resident rows outside these components (the MTP draft + pool's KV rows ride the same block-id space). Under this planner each + component is its own slot, in input order.""" + row_bytes = [c.bytes_per_slot * block_size + c.const_bytes for c in components] + per_block = sum(row_bytes) + reserved_bytes_per_block + num_blocks = budget_bytes // per_block + if num_blocks <= 1: + raise ValueError("budget too small for one usable block") + geo = BlockGeometry( + block_size=block_size, block_bytes=per_block, num_blocks=num_blocks + ) + tensors = tuple( + TensorPlan( + name=f"flat_{c.group_id}_{c.layer}_{c.component}", + nbytes=num_blocks * nbytes, + bindings=(LayerBinding(i, c.group_id, c.layer, c.component, nbytes, 0),), + ) + for i, (c, nbytes) in enumerate(zip(components, row_bytes)) + ) + return FlatMemoryPlan(geometry=geo, tensors=tensors) + + +def plan_tensors(components, *, block_size, alignment, budget_bytes): + """Pair slot j with the j-th layer of every group over one page-id space. + Each slot tensor is sized by its own packed row (the sum of its bindings' + per-block bytes); geometry.block_bytes accounts one block's total across + all slots.""" + geo = solve_page_geometry(components, block_size=block_size, alignment=alignment) + layers_by_group: dict[str, list[int]] = {} + for c in components: + layers = layers_by_group.setdefault(c.group_id, []) + if c.layer not in layers: + layers.append(c.layer) + num_slots = max(len(v) for v in layers_by_group.values()) + + slot_bindings: list[tuple[LayerBinding, ...]] = [] + for slot in range(num_slots): + bindings = [] + for gid, layers in layers_by_group.items(): + if slot >= len(layers): + continue + layer = layers[slot] + row_offset = 0 + for c in components: + if c.group_id != gid or c.layer != layer: + continue + nbytes = c.bytes_per_slot * geo.block_size + c.const_bytes + bindings.append( + LayerBinding(slot, gid, layer, c.component, nbytes, row_offset) + ) + row_offset += nbytes + slot_bindings.append(tuple(bindings)) + slot_rows = [sum(b.nbytes_per_block for b in bs) for bs in slot_bindings] + + num_blocks = budget_bytes // sum(slot_rows) + if num_blocks <= 1: + raise ValueError("budget too small for one usable block per slot") + geo = replace(geo, block_bytes=sum(slot_rows), num_blocks=num_blocks) + + tensors = tuple( + TensorPlan( + name=f"flat_slab_{slot}", + nbytes=num_blocks * slot_rows[slot], + bindings=bindings, + ) + for slot, bindings in enumerate(slot_bindings) + ) + return FlatMemoryPlan(geometry=geo, tensors=tensors) diff --git a/python/tokenspeed/runtime/configs/kimi_k25_config.py b/python/tokenspeed/runtime/configs/kimi_k25_config.py new file mode 100644 index 0000000..ffa9f98 --- /dev/null +++ b/python/tokenspeed/runtime/configs/kimi_k25_config.py @@ -0,0 +1,200 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright 2023-2024 SGLang Team +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +""" +Kimi K25 Model Configuration. +""" + +from transformers import DeepseekV3Config +from transformers.configuration_utils import PretrainedConfig + + +class KimiK25VisionConfig(PretrainedConfig): + """Vision configuration for K2-VL (vision tower + mm projector). + + Args: + Vision Tower Parameters: + patch_size: Patch size for vision tower. + init_pos_emb_height: Initial position embedding height. + init_pos_emb_width: Initial position embedding width. + init_pos_emb_time: Initial position embedding time dimension. + pos_emb_type: Type of position embedding. + num_attention_heads: Number of attention heads in vision tower. + num_hidden_layers: Number of hidden layers in vision tower. + hidden_size: Hidden size of vision tower. + intermediate_size: Intermediate size in vision tower FFN. + merge_kernel_size: Kernel size for spatial patch merging. + video_attn_type: Type of video attention. + merge_type: Type of merge operation. + + MM Projector Parameters: + mm_projector_type: Type of multimodal projector. + mm_hidden_size: Hidden size for projector (defaults to hidden_size). + projector_hidden_act: Activation function for projector. + projector_ln_eps: Layer norm epsilon for projector. + """ + + model_type = "kimi_k25" + + def __init__( + self, + # Vision Tower + patch_size: int = 14, + init_pos_emb_height: int = 64, + init_pos_emb_width: int = 64, + init_pos_emb_time: int = 4, + pos_emb_type: str = "divided_fixed", + 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), + video_attn_type: str = "spatial_temporal", + merge_type: str = "sd2_tpool", + # MM Projector + mm_projector_type: str = "patchmerger", + mm_hidden_size: int | None = None, + projector_hidden_act: str = "gelu", + projector_ln_eps: float = 1e-5, + text_hidden_size: int = 7168, + vt_hidden_size: int | None = None, + **kwargs, + ): + super().__init__(**kwargs) + # Vision Tower + self.patch_size = patch_size + self.init_pos_emb_height = init_pos_emb_height + self.init_pos_emb_width = init_pos_emb_width + self.init_pos_emb_time = init_pos_emb_time + self.pos_emb_type = pos_emb_type + self.num_attention_heads = num_attention_heads + self.num_hidden_layers = num_hidden_layers + self.hidden_size = hidden_size + # Vision-tower hidden size the mm projector reads; defaults to hidden_size. + self.vt_hidden_size = ( + vt_hidden_size if vt_hidden_size is not None else hidden_size + ) + self.intermediate_size = intermediate_size + self.merge_kernel_size = merge_kernel_size + self.video_attn_type = video_attn_type + self.merge_type = merge_type + # MM Projector + self.mm_projector_type = mm_projector_type + if mm_hidden_size is not None: + self.mm_hidden_size = mm_hidden_size + else: + self.mm_hidden_size = hidden_size + self.projector_hidden_act = projector_hidden_act + self.projector_ln_eps = projector_ln_eps + self.text_hidden_size = text_hidden_size + + +class KimiK25Config(PretrainedConfig): + """K2-VL model configuration. + + K2-VL extends Kimi-VL with video support using video-chunks. + A video-chunk consists of multiple consecutive frames (default: 4) + that are processed together with temporal pooling. + + Args: + text_config: Configuration for the text model (DeepseekV3). + + Vision Tower Parameters: + patch_size: Patch size for vision tower. + init_pos_emb_height: Initial position embedding height. + init_pos_emb_width: Initial position embedding width. + init_pos_emb_time: Initial position embedding time dimension. + pos_emb_type: Type of position embedding. + vt_num_attention_heads: Number of attention heads in vision tower. + vt_num_hidden_layers: Number of hidden layers in vision tower. + vt_hidden_size: Hidden size of vision tower. + vt_intermediate_size: Intermediate size in vision tower FFN. + merge_kernel_size: Kernel size for spatial patch merging. + video_attn_type: Type of video attention. + merge_type: Type of merge operation. + + Video-Chunk Parameters: + temporal_merge_kernel_size: Number of frames per video chunk. + Default is 4, meaning 4 frames are merged into 1 chunk. + sample_fps: Video sampling frame rate. + timestamp_mode: Format for chunk timestamps. + + MM Projector Parameters: + mm_projector_type: Type of multimodal projector. + mm_hidden_size: Hidden size from vision tower. + projector_hidden_act: Activation function for projector. + projector_ln_eps: Layer norm epsilon for projector. + + Other Parameters: + ignore_index: The ignore index for the loss function. + media_placeholder_token_id: The token ID for media placeholders. + pad_token_id: The token ID for padding. + """ + + model_type = "kimi_k25" + + def __init__( + self, + text_config: dict | DeepseekV3Config | None = None, + vision_config: dict | KimiK25VisionConfig | None = None, + # Other parameters + ignore_index: int = -100, + media_placeholder_token_id: int = 163605, + pad_token_id: int = 0, + use_unified_vision_chunk: bool = False, + video_placeholder: str = "<|kimi_k25_video_placeholder|>", + **kwargs, + ): + if text_config is None: + text_config = DeepseekV3Config() + elif isinstance(text_config, dict): + text_config = DeepseekV3Config(**text_config) + + if vision_config is None: + vision_config = KimiK25VisionConfig() + elif isinstance(vision_config, dict): + vision_config = KimiK25VisionConfig(**vision_config) + self.vision_config = vision_config + self.text_config = text_config + # Other config + self.ignore_index = ignore_index + self.media_placeholder_token_id = media_placeholder_token_id + self.use_unified_vision_chunk = use_unified_vision_chunk + self.video_placeholder = video_placeholder + + # Propagate quantization config from text model + if getattr(self.text_config, "quantization_config", None) is not None: + self.quantization_config = self.text_config.quantization_config + + super().__init__(pad_token_id=pad_token_id, **kwargs) + + @property + def hidden_size(self) -> int: + """Get hidden size from text config for compatibility.""" + return self.text_config.hidden_size + + @property + def vocab_size(self) -> int: + """Get vocab size from text config for compatibility.""" + return self.text_config.vocab_size diff --git a/python/tokenspeed/runtime/configs/kimi_k2_config.py b/python/tokenspeed/runtime/configs/kimi_k2_config.py new file mode 100644 index 0000000..d2267ff --- /dev/null +++ b/python/tokenspeed/runtime/configs/kimi_k2_config.py @@ -0,0 +1,37 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Kimi K2 model configuration. + +Kimi K2 checkpoints (and their Eagle3 drafts) ship a custom ``model_type`` of +``kimi_k2`` without an ``auto_map`` entry, so ``transformers`` cannot resolve +the config on its own. Structurally the model matches DeepSeek-V3 (MLA + MoE +with ``sigmoid`` / ``noaux_tc`` routing), so we expose a thin subclass that +inherits ``DeepseekV3Config`` and only overrides the registered ``model_type``. +""" + +from transformers import DeepseekV3Config + + +class KimiK2Config(DeepseekV3Config): + model_type = "kimi_k2" + + +__all__ = ["KimiK2Config"] diff --git a/python/tokenspeed/runtime/configs/load_config.py b/python/tokenspeed/runtime/configs/load_config.py new file mode 100755 index 0000000..c443fb9 --- /dev/null +++ b/python/tokenspeed/runtime/configs/load_config.py @@ -0,0 +1,94 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Model weight loading configuration.""" + +import enum +import json +from dataclasses import dataclass, field + +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +class LoadFormat(str, enum.Enum): + AUTO = "auto" + PT = "pt" + SAFETENSORS = "safetensors" + NPCACHE = "npcache" + DUMMY = "dummy" + SHARDED_STATE = "sharded_state" + MISTRAL = "mistral" + EXTENSIBLE = "extensible" + + +@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. + 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: str | LoadFormat = LoadFormat.AUTO + download_dir: str | None = None + model_loader_extra_config: str | dict | None = field(default_factory=dict) + ignore_patterns: list[str] | str | None = None + decryption_key_file: str | None = None + weight_loader_prefetch_checkpoints: bool = False + weight_loader_prefetch_num_threads: int = 4 + + ext_yaml: str | None = None + + def __post_init__(self) -> None: + 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) diff --git a/python/tokenspeed/runtime/configs/minimax_m2_config.py b/python/tokenspeed/runtime/configs/minimax_m2_config.py new file mode 100644 index 0000000..e513e20 --- /dev/null +++ b/python/tokenspeed/runtime/configs/minimax_m2_config.py @@ -0,0 +1,127 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""MiniMax-M2 model configuration definitions.""" + +from transformers.configuration_utils import PretrainedConfig +from transformers.utils import logging + +from tokenspeed.runtime.configs.utils import rope_config_validation + +logger = logging.get_logger(__name__) + + +class MiniMaxM2Config(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`MiniMaxM2Model`]. + It is used to instantiate MiniMax-M2 family models according to the specified arguments, + defining the model architecture. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. + """ + + model_type = "minimax_m2" + keys_to_ignore_at_inference = ["past_key_values"] + + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.block_sparse_moe.gate": "colwise_rep", + "layers.*.block_sparse_moe.experts.*.w1": "colwise", + "layers.*.block_sparse_moe.experts.*.w2": "rowwise", + "layers.*.block_sparse_moe.experts.*.w3": "colwise", + } + + def __init__( + self, + vocab_size=200064, + hidden_size=3072, + intermediate_size=1536, + num_hidden_layers=62, + num_attention_heads=48, + num_key_value_heads=8, + head_dim=128, + hidden_act="silu", + max_position_embeddings=196608, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + tie_word_embeddings=False, + rope_theta=5_000_000, + rope_scaling=None, + rotary_dim=64, + attention_bias=False, + attention_dropout=0.0, + # MoE + num_local_experts=256, + num_experts_per_tok=8, + scoring_func="sigmoid", + use_routing_bias=True, + norm_topk_prob=False, + output_router_logits=False, + router_aux_loss_coef=0.001, + # QK-Norm + use_qk_norm=True, + qk_norm_type="per_layer", + **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.head_dim = head_dim + 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.rotary_dim = rotary_dim + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + # Validate rope + if self.rope_scaling is not None and "type" in self.rope_scaling: + self.rope_scaling["rope_type"] = self.rope_scaling["type"] + rope_config_validation(self) + # MoE + self.num_local_experts = num_local_experts + self.num_experts_per_tok = num_experts_per_tok + self.scoring_func = scoring_func + self.use_routing_bias = use_routing_bias + self.norm_topk_prob = norm_topk_prob + self.output_router_logits = output_router_logits + self.router_aux_loss_coef = router_aux_loss_coef + # QK-Norm + self.use_qk_norm = use_qk_norm + self.qk_norm_type = qk_norm_type + # Preserve extra public checkpoint metadata through PretrainedConfig + # without making it part of the MiniMax-M2 serving runtime. + super().__init__( + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + +__all__ = ["MiniMaxM2Config"] diff --git a/python/tokenspeed/runtime/configs/model_config.py b/python/tokenspeed/runtime/configs/model_config.py new file mode 100644 index 0000000..42a5d78 --- /dev/null +++ b/python/tokenspeed/runtime/configs/model_config.py @@ -0,0 +1,738 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Model configuration helpers and derived runtime metadata.""" + +import copy +import json +import math +import os +from collections.abc import Callable +from dataclasses import dataclass +from enum import IntEnum, auto + +import torch +import yaml +from transformers import PretrainedConfig + +from tokenspeed.runtime.layers.quantization import QUANTIZATION_METHODS +from tokenspeed.runtime.utils import get_colorful_logger +from tokenspeed.runtime.utils.env import envs +from tokenspeed.runtime.utils.hf_transformers_utils import ( + get_config, + get_context_length, + get_generation_config, + resolve_architecture, +) +from tokenspeed.runtime.utils.server_args import ServerArgs + +logger = get_colorful_logger(__name__) + +_DEEPSEEK_V4_ARCHITECTURES = frozenset( + { + "DeepseekV4ForCausalLM", + "DeepseekV4ForCausalLMNextN", + } +) +_MLA_ARCHITECTURES = frozenset( + { + "DeepseekV3ForCausalLM", + "DeepseekV3ForCausalLMNextN", + "Eagle3DeepseekV2ForCausalLM", + "LongcatFlashForCausalLM", + "KimiK25ForConditionalGeneration", + } +) +_DSA_ARCHITECTURES = frozenset( + { + "GlmMoeDsaForCausalLM", + "GlmMoeDsaForCausalLMNextN", + } +) +_DOUBLE_ATTENTION_LAYER_ARCHITECTURES = frozenset( + { + "LongcatFlashForCausalLM", + } +) + + +class AttentionArch(IntEnum): + MLA = auto() + MHA = auto() + DSA = auto() + + +@dataclass(frozen=True) +class _AttentionFamilySpec: + name: str + architectures: frozenset[str] + configure: Callable[[object], None] + default_backend: str | None = None + supports_target_verify_forward_mode: bool = False + default_block_size: int | None = None + + +def override_model_config(model_config, ext_yaml): + with open(ext_yaml, encoding="utf-8") as f: + ext_config = yaml.safe_load(f) + + override_model_config: dict = ext_config.get("override_model_config", {}) + for k, v in override_model_config.items(): + if hasattr(model_config, k): + old_v = model_config.__getattribute__(k) + if isinstance(v, dict): + new_v = copy.deepcopy(old_v) + new_v.__dict__.update(v) + else: + new_v = v + model_config.__setattr__(k, new_v) + logger.info("Override model config: %s=%r", k, new_v) + + +def is_deepseek_v4(config: PretrainedConfig) -> bool: + return resolve_architecture(config) in _DEEPSEEK_V4_ARCHITECTURES + + +def is_deepseek_v4_nextn(config: PretrainedConfig) -> bool: + return resolve_architecture(config) == "DeepseekV4ForCausalLMNextN" + + +def configure_deepseek_v4_attention(model_config) -> None: + """Derive DeepSeek V4's MLA-like dimensions for runtime setup.""" + + hf_config = model_config.hf_config + model_config.head_dim = hf_config.head_dim + model_config.attention_arch = AttentionArch.MLA + model_config.kv_lora_rank = hf_config.head_dim + model_config.qk_rope_head_dim = hf_config.qk_rope_head_dim + model_config.qk_nope_head_dim = hf_config.head_dim - hf_config.qk_rope_head_dim + model_config.v_head_dim = hf_config.head_dim + model_config.index_head_dim = getattr(hf_config, "index_head_dim", None) + model_config.scaling = 1 / math.sqrt(model_config.head_dim) + rope_scaling = getattr(hf_config, "rope_scaling", None) + if rope_scaling: + mscale_all_dim = rope_scaling.get("mscale_all_dim", False) + scaling_factor = rope_scaling["factor"] + mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim)) + model_config.scaling = model_config.scaling * mscale * mscale + + +def configure_glm_attention(model_config) -> None: + mla_config = ( + model_config.hf_text_config + if hasattr(model_config.hf_text_config, "kv_lora_rank") + else model_config.hf_config + ) + required_fields = ( + "kv_lora_rank", + "qk_nope_head_dim", + "qk_rope_head_dim", + "v_head_dim", + "index_topk", + "index_head_dim", + "index_n_heads", + ) + missing_fields = [ + field for field in required_fields if not hasattr(mla_config, field) + ] + if missing_fields: + raise ValueError( + "GLM attention config is missing required fields: " + + ", ".join(missing_fields) + ) + + model_config.head_dim = getattr(mla_config, "qk_head_dim", None) + if model_config.head_dim is None: + model_config.head_dim = ( + mla_config.qk_nope_head_dim + mla_config.qk_rope_head_dim + ) + model_config.attention_arch = AttentionArch.DSA + model_config.kv_lora_rank = mla_config.kv_lora_rank + model_config.qk_nope_head_dim = mla_config.qk_nope_head_dim + model_config.qk_rope_head_dim = mla_config.qk_rope_head_dim + model_config.v_head_dim = mla_config.v_head_dim + model_config.index_topk = mla_config.index_topk + model_config.index_head_dim = mla_config.index_head_dim + model_config.index_n_heads = mla_config.index_n_heads + model_config.index_topk_pattern = getattr(mla_config, "index_topk_pattern", None) + + model_config.scaling = 1 / math.sqrt( + model_config.qk_nope_head_dim + model_config.qk_rope_head_dim + ) + rope_scaling = getattr(mla_config, "rope_scaling", None) + if rope_scaling and "factor" in rope_scaling: + mscale_all_dim = rope_scaling.get("mscale_all_dim", False) + scaling_factor = rope_scaling["factor"] + mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim)) + model_config.scaling = model_config.scaling * mscale * mscale + + +def configure_mla_attention(model_config) -> None: + mla_config = ( + model_config.hf_text_config + if hasattr(model_config.hf_text_config, "kv_lora_rank") + else model_config.hf_config + ) + model_config.head_dim = 256 + model_config.attention_arch = AttentionArch.MLA + model_config.kv_lora_rank = mla_config.kv_lora_rank + model_config.qk_nope_head_dim = mla_config.qk_nope_head_dim + model_config.qk_rope_head_dim = mla_config.qk_rope_head_dim + model_config.v_head_dim = mla_config.v_head_dim + + model_config.scaling = 1 / math.sqrt( + model_config.qk_nope_head_dim + model_config.qk_rope_head_dim + ) + rope_scaling = getattr(mla_config, "rope_scaling", None) + if rope_scaling and "factor" in rope_scaling: + mscale_all_dim = rope_scaling.get("mscale_all_dim", False) + scaling_factor = rope_scaling["factor"] + mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim)) + model_config.scaling = model_config.scaling * mscale * mscale + + +_ATTENTION_FAMILY_SPECS = ( + _AttentionFamilySpec( + name="DeepSeek V4", + architectures=_DEEPSEEK_V4_ARCHITECTURES, + configure=configure_deepseek_v4_attention, + supports_target_verify_forward_mode=True, + default_block_size=256, + ), + _AttentionFamilySpec( + name="GLM", + architectures=_DSA_ARCHITECTURES, + configure=configure_glm_attention, + default_backend="dsa", + supports_target_verify_forward_mode=True, + ), + _AttentionFamilySpec( + name="MLA", + architectures=_MLA_ARCHITECTURES, + configure=configure_mla_attention, + ), +) + + +def _model_architectures( + hf_config: PretrainedConfig, + hf_text_config: PretrainedConfig, +) -> list[str]: + return ( + [resolve_architecture(hf_config)] + + list(getattr(hf_config, "architectures", None) or []) + + list(getattr(hf_text_config, "architectures", []) or []) + ) + + +def _resolve_attention_family( + hf_config: PretrainedConfig, + hf_text_config: PretrainedConfig, +) -> _AttentionFamilySpec | None: + architectures = _model_architectures(hf_config, hf_text_config) + for spec in _ATTENTION_FAMILY_SPECS: + if any(arch in spec.architectures for arch in architectures): + return spec + return None + + +def _apply_attention_family_defaults( + server_args: ServerArgs, + spec: _AttentionFamilySpec, +) -> None: + if spec.default_block_size is not None: + block_size_default = ServerArgs.__dataclass_fields__["block_size"].default + if server_args.block_size == block_size_default: + logger.info( + "%s default block_size=%d; pass --block-size with a value other " + "than %d to keep that value.", + spec.name, + spec.default_block_size, + block_size_default, + ) + server_args.block_size = spec.default_block_size + if spec.default_backend is not None and server_args.attention_backend is None: + server_args.attention_backend = spec.default_backend + + +def _derive_num_attention_layers( + hf_config: PretrainedConfig, + num_hidden_layers: int, +) -> int: + architectures = getattr(hf_config, "architectures", None) or [] + num_attention_layers = num_hidden_layers + if is_deepseek_v4_nextn(hf_config): + num_attention_layers = int(getattr(hf_config, "num_nextn_predict_layers", 1)) + if any(arch in _DOUBLE_ATTENTION_LAYER_ARCHITECTURES for arch in architectures): + num_attention_layers = num_hidden_layers * 2 + return num_attention_layers + + +class ModelConfig: + def __init__( + self, + model_path: str, + trust_remote_code: bool = True, + revision: str | None = None, + context_length: int | None = None, + model_override_args: dict | None = None, + dtype: str = "auto", + quantization: str | None = None, + override_config_file: str | None = None, + is_draft_worker: bool | None = False, + server_args: ServerArgs = None, + ) -> None: + self.model_path = model_path + self.revision = revision + self.quantization = quantization + self.mapping = server_args.mapping + + # Parse args + 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( + model_path, + trust_remote_code=trust_remote_code, + revision=revision, + model_override_args=self.model_override_args, + is_draft_worker=is_draft_worker, + **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) + + # Check model type + self.is_generation = is_generation_model(self.hf_config.architectures) + self.is_multimodal = is_multimodal_model(self.hf_config.architectures) + self.is_multimodal_gen = is_multimodal_gen_model(self.hf_config.architectures) + self.is_image_gen = is_image_gen_model(self.hf_config.architectures) + self.is_audio_model = is_audio_model(self.hf_config.architectures) + + language_model_only = bool(getattr(server_args, "language_model_only", False)) + # Target-only flag; never apply to draft / auxiliary checkpoints. + apply_language_model_only = language_model_only and not is_draft_worker + if apply_language_model_only: + if not self.is_multimodal: + raise ValueError( + "--language-model-only requires a multimodal model checkpoint." + ) + logger.info( + "Running in language-model-only mode: vision/audio encoders will " + "be skipped; requests with multimodal inputs will be rejected." + ) + # ``is_multimodal`` is the architectural fact; this is the runtime gate. + self.is_multimodal_active = self.is_multimodal and not apply_language_model_only + # Vision-only role (EPD encode): the inverse axis of language_model_only. + # Build the vision tower (is_multimodal_active stays True) but SKIP LM + # construction + LM weight load so a full ViT fits at encode TP=1. + encoder_only = ( + getattr(server_args, "disaggregation_mode", None) == "encode" + and not is_draft_worker + ) + if encoder_only and not self.is_multimodal: + raise ValueError( + "disaggregation_mode=encode requires a multimodal checkpoint." + ) + if encoder_only and apply_language_model_only: + raise ValueError( + "disaggregation_mode=encode (encoder-only) and language_model_only " + "are mutually exclusive." + ) + if encoder_only and self.is_audio_model: + raise ValueError( + "disaggregation_mode=encode does not support audio models; " + "only image/video encoders are currently supported." + ) + if encoder_only: + # Single model-facing gate: Kimi reads hf_config.encoder_only directly; + # Qwen3_5ForConditionalGeneration reads it to skip LM construction. + self.hf_config.encoder_only = True + logger.info( + "Running in encoder-only mode: the language model will not " + "be constructed or loaded (encode role)." + ) + # Cap gpu_memory_utilization for VLMs in mm mode — the vision encoder + # needs headroom that the global default doesn't account for. + if ( + self.is_multimodal_active + and getattr(server_args, "_gpu_memory_utilization_defaulted", False) + and server_args.gpu_memory_utilization > 0.9 + ): + logger.info( + "Clamping gpu_memory_utilization %.2f -> 0.9 to leave headroom " + "for the vision encoder.", + server_args.gpu_memory_utilization, + ) + server_args.gpu_memory_utilization = 0.9 + self.mm_attention_backend = getattr(server_args, "mm_attention_backend", None) + 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: + if envs.TOKENSPEED_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN.get(): + logger.warning( + "User-specified context_length (%s) is greater than the derived " + "context_length (%s). This may lead to incorrect model outputs or " + "CUDA errors.", + context_length, + derived_context_len, + ) + self.context_len = context_length + else: + raise ValueError( + f"User-specified 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. " + f"To allow overriding this maximum, set the env var TOKENSPEED_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, + ) + + # MLA/DSA families carry per-head dimension metadata that does not + # follow the standard hidden_size / num_attention_heads derivation above. + attention_family = _resolve_attention_family( + self.hf_config, + self.hf_text_config, + ) + if attention_family is not None: + _apply_attention_family_defaults(server_args, attention_family) + attention_family.configure(self) + 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 + else: + self.attention_arch = AttentionArch.MHA + + self.use_v4_mtp_paged_metadata = ( + getattr(server_args, "speculative_algorithm", None) is not None + and not is_draft_worker + and attention_family is not None + and attention_family.supports_target_verify_forward_mode + ) + + 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 = getattr(self.hf_text_config, "num_hidden_layers", None) + if self.num_hidden_layers is None: + self.num_hidden_layers = self.hf_text_config.num_layers + self.num_attention_layers = _derive_num_attention_layers( + self.hf_config, + self.num_hidden_layers, + ) + if is_draft_worker: + mtp_layers = getattr(self.hf_text_config, "mtp_num_hidden_layers", None) + if mtp_layers is not None: + self.num_attention_layers = mtp_layers + else: + nextn_layers = getattr( + self.hf_text_config, "num_nextn_predict_layers", None + ) + if nextn_layers is not None and nextn_layers > 0: + self.num_attention_layers = nextn_layers + self.vocab_size = self.hf_text_config.vocab_size + + # Verify quantization + self._verify_quantization() + + # Cache attributes + self.hf_eos_token_id = self.get_hf_eos_token_id() + self.image_token_id = getattr(self.hf_config, "image_token_id", None) + + if server_args is not None and server_args.load_format == "extensible": + override_model_config(self, server_args.ext_yaml) + + 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: + # modelopt NVFP4 checkpoints store quant config in hf_quant_config.json + # Resolve the local snapshot directory (model_path may be a HF hub ID) + if os.path.isdir(self.model_path): + model_dir = self.model_path + else: + try: + from huggingface_hub import snapshot_download + + model_dir = snapshot_download( + self.model_path, + revision=self.revision, + allow_patterns=["*.json"], + local_files_only=True, + ) + except Exception as exc: + logger.debug( + "Unable to resolve local quantization config for %s: %s", + self.model_path, + exc, + ) + model_dir = None + if model_dir is not None: + hf_quant_path = os.path.join(model_dir, "hf_quant_config.json") + if os.path.isfile(hf_quant_path): + with open(hf_quant_path, encoding="utf-8") as f: + hf_quant = json.load(f) + quant_algo = hf_quant.get("quantization", {}).get("quant_algo", "") + if quant_algo: + quant_cfg = { + "quant_method": "modelopt", + "quant_algo": quant_algo, + } + quant_cfg.update(hf_quant.get("quantization", {})) + return quant_cfg + + def _verify_quantization(self) -> None: + supported_quantization = [*QUANTIZATION_METHODS] + + optimized_quantization_methods = [ + "fp8", + "nvfp4", + "mxfp4", + "compressed_tensors", + "compressed-tensors", + "w8a8_fp8", + ] + compatible_quantization_methods = { + "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", "").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 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 get_hf_eos_token_id(self) -> set[int] | None: + eos_ids = getattr(self.hf_config, "eos_token_id", None) + if eos_ids: + # 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 get_hf_text_config(config: PretrainedConfig): + """Get the "sub" config relevant to llm for multi modal models. + No op for pure text models. + """ + class_name = resolve_architecture(config) + 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. + # We set `dtype` of config to `torch.float16` for the weights, as + # `torch.float16` is default used for image features in + # `python/tokenspeed/runtime/models/llava.py`. + config.dtype = torch.float16 + return config + + if hasattr(config, "thinker_config"): + thinker_config = config.thinker_config + if hasattr(thinker_config, "text_config"): + return thinker_config.text_config + return thinker_config + if hasattr(config, "text_config"): + if not hasattr(config.text_config, "num_attention_heads"): + raise ValueError("text_config must define num_attention_heads") + return config.text_config + return config + + +_STR_DTYPE_TO_TORCH_DTYPE = { + "half": torch.float16, + "float16": torch.float16, + "float": torch.float32, + "float32": torch.float32, + "bfloat16": torch.bfloat16, +} + + +def _get_and_verify_dtype( + config: PretrainedConfig, + dtype: str | torch.dtype, +) -> torch.dtype: + # config.dtype can be missing or None. + config_dtype = getattr(config, "dtype", None) + if config_dtype is None: + config_dtype = torch.bfloat16 + + if isinstance(dtype, str): + dtype = dtype.lower() + if dtype == "auto": + if config_dtype == torch.float32: + if config.model_type == "gemma2": + logger.info( + "For Gemma 2, 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) + elif config_dtype == torch.float32: + # Downcasting from float32 to float16 or bfloat16 is allowed. + logger.info("Downcasting %s to %s.", config_dtype, torch_dtype) + 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]): + return True + + +def is_multimodal_model(model_architectures: list[str] | None): + multimodal_architectures = { + "Qwen3_5ForConditionalGeneration", + "Qwen3_5MoeForConditionalGeneration", + "Qwen3OmniMoeForConditionalGeneration", + "Qwen3ASRForConditionalGeneration", + "KimiK25ForConditionalGeneration", + } + return any(arch in multimodal_architectures for arch in model_architectures or []) + + +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] | None): + audio_architectures = { + "Qwen3OmniMoeForConditionalGeneration", + "Qwen3ASRForConditionalGeneration", + } + return any(arch in audio_architectures for arch in model_architectures or []) + + +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 diff --git a/python/tokenspeed/runtime/configs/paged_cache_spec.py b/python/tokenspeed/runtime/configs/paged_cache_spec.py new file mode 100644 index 0000000..17d0a4b --- /dev/null +++ b/python/tokenspeed/runtime/configs/paged_cache_spec.py @@ -0,0 +1,573 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Literal + +Retention = Literal["full_history", "sliding_window"] +Family = Literal["history", "state"] + + +@dataclass(frozen=True) +class PagedCacheGroupSpec: + group_id: str + retention: Retention + rows_per_page: int + entry_stride_tokens: int + sliding_window_tokens: int | None + # History groups form a chain; State groups only need the trailing window. + family: Family = "history" + + +_PAGED_CACHE_GROUP_DUMMY_PAGES = 1 + + +def scheduler_ext_flat_kvcache() -> bool: + """True iff the installed tokenspeed_scheduler ext was built with + TOKENSPEED_FLAT_KVCACHE. A missing package or an older / radix-built + ext reports False — the radix-safe default (never delivers flat tables). + """ + try: + # Local import: module must stay importable without the compiled ext. + import tokenspeed_scheduler + except ImportError: + return False + return bool(getattr(tokenspeed_scheduler, "FLAT_KVCACHE", False)) + + +# Paged-cache label vocabulary (NOT the HF checkpoint's serialized enum: +# Qwen3.5 checkpoints spell full attention "attention"). +FULL_ATTENTION = "full_attention" +LINEAR_ATTENTION = "linear_attention" + +# layer_type label -> retention. GPT-OSS uses the first two, Qwen3.5 GDN +# layers use "linear_attention"; unknown labels raise. +_LAYER_TYPE_RETENTION: dict[str, Retention] = { + FULL_ATTENTION: "full_history", + "sliding_attention": "sliding_window", + # State groups ride full_history retention: the C++ side keys the + # mamba-state kind on family == State && retention != SlidingWindow. + LINEAR_ATTENTION: "full_history", +} + +# Labels whose group is state-family (recurrent state rows, not KV history). +STATE_LAYER_TYPES = frozenset({LINEAR_ATTENTION}) + + +def hybrid_slab_group_size( + layer_types: Sequence[str] | None, + *, + sliding_window_tokens: int | Sequence[int | None] | None = None, +) -> int | None: + """Group size for the hybrid slab KV layout (one layer of EACH group + shares a K/V slab), or None to keep the legacy per-layer layout. + + Single source (canonical) for both the sizing divisor (registry KV + profile) and the buffer layout (_create_buffers) -- the two must never + disagree. Safe only with the flat ext (its single BlockPool owns each + page id by at most one group, so paired layers' live rows never + overlap) and equal group sizes. Unknown labels degrade to None -- the + predicate gates an optimization, so it must not raise. + + Multi-window models (a per-layer window sequence with >1 distinct + window) degrade to None: the slab pairing is per raw label, not per + (retention, window) group. + """ + if not scheduler_ext_flat_kvcache(): + return None + if not layer_types: + return None + counts: dict[str, int] = {} + for label in layer_types: + # State rows are not byte-equal with KV rows, so no slab pairing. + if label not in _LAYER_TYPE_RETENTION or label in STATE_LAYER_TYPES: + return None + counts[label] = counts.get(label, 0) + 1 + if len(counts) < 2: + return None + if sliding_window_tokens is not None and not isinstance(sliding_window_tokens, int): + if not isinstance(sliding_window_tokens, Sequence) or len( + sliding_window_tokens + ) != len(layer_types): + return None + distinct = { + w + for label, w in zip(layer_types, sliding_window_tokens) + if _LAYER_TYPE_RETENTION[label] == "sliding_window" + and isinstance(w, int) + and not isinstance(w, bool) + and w > 0 + } + if len(distinct) > 1: + return None + sizes = set(counts.values()) + if len(sizes) != 1: + return None + return sizes.pop() + + +def _flat_kv_backend(attn_backend: object) -> object: + """The backend whose KV-table consumption matters for flat safety: the + backend itself, or a composite's full-attention sub-backend (hybrid's + per-layer KV routing lives there and is user-selectable). The linear + side consumes only the state group's table through its own explicit + flat path and is out of scope here. + """ + sub = getattr(attn_backend, "full_attn_backend", None) + if sub is not None: + return _flat_kv_backend(sub) + return attn_backend + + +def validate_flat_scheduler_config( + *, + flat_kvcache_ext: bool, + paged_cache_groups: Sequence[object], + attn_backend: object, + kv_pool: object, + speculative_algorithm: str | None = None, +) -> None: + """Fail fast, before the C++ ``Scheduler`` ctor, when a flat-built ext + cannot drive this setup: a backend (or hybrid sub-backend) that would not + consume the per-group flat tables, or zero published groups. No-op on a + radix build. + """ + if not flat_kvcache_ext: + return + pool_name = type(kv_pool).__name__ + backend = _flat_kv_backend(attn_backend) + backend_name = type(backend).__name__ + uses_paged = bool(getattr(backend, "uses_paged_cache_groups", False)) + uses_flat = bool(getattr(backend, "uses_flat_cache_groups", False)) + if uses_paged and not uses_flat: + raise RuntimeError( + "flat scheduler build (TOKENSPEED_FLAT_KVCACHE) does not support " + f"this model's cache layout yet: attention backend {backend_name} " + f"(KV pool {pool_name}) consumes paged-cache groups through the " + "radix scheduler's populate path, which the flat build compiles " + "out — CUDA graphs would silently replay against stale capture " + "placeholders. Use a radix-built tokenspeed_scheduler extension " + "for this model." + ) + if speculative_algorithm is not None and not getattr( + backend, "flat_spec_capable", True + ): + raise RuntimeError( + "flat scheduler build (TOKENSPEED_FLAT_KVCACHE): attention backend " + f"{backend_name} does not support flat cache groups with " + "speculative decoding yet. Use the MHA backend or a radix-built " + "tokenspeed_scheduler extension." + ) + if speculative_algorithm == "DFLASH": + raise RuntimeError( + "flat scheduler build (TOKENSPEED_FLAT_KVCACHE): DFLASH block " + "decode is unsupported on the flat path. Use a radix-built " + "tokenspeed_scheduler extension." + ) + if len(paged_cache_groups) > 1 and not uses_flat: + # A table-blind backend on a multi-group pool would index every + # layer through the C++ single-table fallback (a first-group + # sample) — with slab-aliased layouts that silently corrupts KV + # past the sliding window. Refuse at startup instead. + raise RuntimeError( + "flat scheduler build (TOKENSPEED_FLAT_KVCACHE): KV pool " + f"{pool_name} publishes {len(paged_cache_groups)} cache groups " + f"but attention backend {backend_name} does not consume flat " + "per-group tables (uses_flat_cache_groups=False); the single-" + "table fallback would serve one group's pages to every layer, " + "silently corrupting KV. Pick a flat-capable attention backend " + "or use a radix-built tokenspeed_scheduler extension." + ) + if not paged_cache_groups: + raise RuntimeError( + "flat scheduler build (TOKENSPEED_FLAT_KVCACHE) requires at least " + f"one paged-cache group, but KV pool {pool_name} publishes none " + "(e.g. mamba/state-only pools). Use a radix-built " + "tokenspeed_scheduler extension for this model." + ) + + +def compute_paged_cache_group_page_counts( + specs: Sequence[PagedCacheGroupSpec], + *, + max_live_requests: int, + max_scheduled_tokens: int, + max_total_tokens: int, + max_context_len: int, + decode_input_tokens: int = 1, + overlap_schedule_depth: int = 0, + safety_margin: int = 0, +) -> dict[str, int]: + # Local import: keeps this module torch-free at import time. + from tokenspeed.runtime.utils.common import ceil_div + + if max_live_requests < 0: + raise ValueError(f"max_live_requests must be >= 0, got {max_live_requests}") + if max_scheduled_tokens < 0: + raise ValueError( + f"max_scheduled_tokens must be >= 0, got {max_scheduled_tokens}" + ) + if max_total_tokens < 0: + raise ValueError(f"max_total_tokens must be >= 0, got {max_total_tokens}") + if max_context_len < 0: + raise ValueError(f"max_context_len must be >= 0, got {max_context_len}") + if decode_input_tokens < 0: + raise ValueError(f"decode_input_tokens must be >= 0, got {decode_input_tokens}") + if overlap_schedule_depth not in (0, 1): + raise ValueError( + f"overlap_schedule_depth must be 0 or 1, got {overlap_schedule_depth}" + ) + if overlap_schedule_depth > 0 and decode_input_tokens == 0: + raise ValueError( + "overlapped paged-cache sizing requires decode_input_tokens > 0" + ) + if safety_margin < 0: + raise ValueError(f"safety_margin must be >= 0, got {safety_margin}") + + counts: dict[str, int] = {} + for spec in specs: + raw_per_page = spec.rows_per_page * spec.entry_stride_tokens + if raw_per_page <= 0: + raise ValueError( + f"PagedCacheGroupSpec {spec.group_id}: rows_per_page * " + "entry_stride_tokens must be > 0" + ) + protected_pages = max_live_requests * ceil_div( + overlap_schedule_depth * decode_input_tokens, raw_per_page + ) + # Mamba-state kind = family "state" AND retention != sliding_window + # (the C++ side keys it the same way); V4's sliding-window state tail + # buffers keep the sliding-window formula below. + if spec.family == "state" and spec.retention == "full_history": + # State group: 2 live pages/request (the W=2 write window) + + # floor(T/P) snapshot pages (snapshots are bounded by the shared + # page-id space), capped at the full-history count. + full_history_total = ( + ceil_div(max_total_tokens, raw_per_page) + + max_live_requests + + protected_pages + + _PAGED_CACHE_GROUP_DUMMY_PAGES + + safety_margin + ) + state_total = ( + max_live_requests * 2 + + max_total_tokens // raw_per_page + + protected_pages + + _PAGED_CACHE_GROUP_DUMMY_PAGES + + safety_margin + ) + total = min(state_total, full_history_total) + elif spec.retention == "full_history": + full_pages = ceil_div(max_total_tokens, raw_per_page) + total = ( + full_pages + + max_live_requests + + protected_pages + + _PAGED_CACHE_GROUP_DUMMY_PAGES + + safety_margin + ) + elif spec.retention == "sliding_window": + window = spec.sliding_window_tokens + if window is None or window <= 0: + raise ValueError( + f"PagedCacheGroupSpec {spec.group_id}: sliding group missing " + "positive sliding_window_tokens" + ) + # Capacity tracks resident history before the next token. + resident_tokens_per_req = min(max(window - 1, 0), max_context_len) + resident_pages = max_live_requests * ceil_div( + resident_tokens_per_req, raw_per_page + ) + scheduled_tokens = min(max_scheduled_tokens, max_total_tokens) + scheduled_pages = ceil_div(scheduled_tokens, raw_per_page) + total = ( + resident_pages + + scheduled_pages + + max_live_requests + + protected_pages + + _PAGED_CACHE_GROUP_DUMMY_PAGES + + safety_margin + ) + else: + raise ValueError( + f"PagedCacheGroupSpec {spec.group_id}: unsupported retention " + f"{spec.retention!r}" + ) + counts[spec.group_id] = int(total) + return counts + + +def _layer_specs( + layer_types: Sequence[str], + sliding_window_tokens: int | Sequence[int | None] | None, +) -> list[tuple[str, Retention, int | None]]: + """Per-layer (group_id, retention, window). group_id is the bare label + unless sliding layers carry more than one distinct window (then + label_), so single-window models keep byte-identical ids. + A scalar window broadcasts to sliding layers; a sequence lines up 1:1.""" + if isinstance(sliding_window_tokens, str): + raise ValueError( + "_layer_specs: sliding_window_tokens must be None, an int, or a " + f"sequence of int/None, got {sliding_window_tokens!r}" + ) + if sliding_window_tokens is None or isinstance(sliding_window_tokens, int): + if isinstance(sliding_window_tokens, bool): + raise ValueError( + "_layer_specs: sliding_window_tokens must be None, an int, or " + f"a sequence of int/None, got {sliding_window_tokens!r}" + ) + windows: list[int | None] = [sliding_window_tokens] * len(layer_types) + scalar = True + elif not isinstance(sliding_window_tokens, Sequence): + raise ValueError( + "_layer_specs: sliding_window_tokens must be None, an int, or a " + f"sequence of int/None, got {sliding_window_tokens!r}" + ) + else: + windows = list(sliding_window_tokens) + scalar = False + if len(windows) != len(layer_types): + raise ValueError( + f"_layer_specs: sliding_window_tokens has {len(windows)} " + f"entries but layer_types has {len(layer_types)}" + ) + rows: list[tuple[str, Retention, int | None]] = [] + for i, (label, raw) in enumerate(zip(layer_types, windows)): + retention = _LAYER_TYPE_RETENTION.get(label) + if retention is None: + raise ValueError( + f"_layer_specs: unknown layer_type {label!r} at layer {i}; " + f"expected one of {sorted(_LAYER_TYPE_RETENTION)}" + ) + if raw is not None and (isinstance(raw, bool) or not isinstance(raw, int)): + raise ValueError( + f"_layer_specs: layer {i} ({label!r}) window must be None or " + f"an int, got {raw!r}" + ) + window = raw + if retention == "sliding_window": + if window is None or window <= 0: + raise ValueError( + f"_layer_specs: layer {i} ({label!r}) is sliding but its " + f"window is not a positive int (got {raw!r})" + ) + else: + if not scalar and window is not None and window > 0: + raise ValueError( + f"_layer_specs: layer {i} ({label!r}) is full-history but " + f"carries sliding window {window}; mislabeled layer_type?" + ) + window = None + rows.append((label, retention, window)) + distinct = {w for _, r, w in rows if r == "sliding_window"} + multi_window = len(distinct) > 1 + return [ + ( + ( + f"{label}_{window}" + if multi_window and retention == "sliding_window" + else label + ), + retention, + window, + ) + for label, retention, window in rows + ] + + +def layer_group_ids( + *, + layer_types: Sequence[str], + sliding_window_tokens: int | Sequence[int | None] | None, +) -> list[str]: + """Per-layer paged-cache group id — the single source multi-window models + will assign ``PagedAttention(group_id=...)`` from (today gpt_oss.py + assigns group_id=layer_type, identical in the single-window case), so + ``flat_block_tables`` keys line up with the published group specs.""" + return [gid for gid, _, _ in _layer_specs(layer_types, sliding_window_tokens)] + + +def group_specs_from_layer_types( + *, + layer_types: Sequence[str], + sliding_window_tokens: int | Sequence[int | None] | None, + page_size: int, +) -> list[PagedCacheGroupSpec]: + """Derive paged-cache group specs from per-layer attention types. + + vLLM-style spec-value grouping: layers collapse into one group per + distinct (retention, window). Group order = first-appearance order. + + Args: + layer_types: Per-layer labels: "full_attention" / "sliding_attention" + / "linear_attention" (state-family, e.g. Qwen3.5 GDN). + sliding_window_tokens: One window for all sliding layers (today's HF + scalar), or a per-layer sequence (multi-window models; full-layer + positions must be None). + page_size: Tokens per page (uniform across groups). + + Raises: + ValueError: unknown label; window sequence length mismatch; sliding + layer without a positive window; full layer carrying a window. + """ + specs: list[PagedCacheGroupSpec] = [] + seen: set[str] = set() + for gid, retention, window in _layer_specs(layer_types, sliding_window_tokens): + if gid in seen: + continue + seen.add(gid) + specs.append( + PagedCacheGroupSpec( + group_id=gid, + retention=retention, + rows_per_page=page_size, + entry_stride_tokens=1, + sliding_window_tokens=window, + family="state" if gid in STATE_LAYER_TYPES else "history", + ) + ) + return specs + + +def publish_paged_cache_groups( + *, + layer_types: Sequence[str], + sliding_window_tokens: int | Sequence[int | None] | None, + page_size: int, + max_live_requests: int, + max_scheduled_tokens: int, + max_total_tokens: int, + max_context_len: int, +) -> tuple[list[PagedCacheGroupSpec], dict[str, int]] | None: + """Publication rule (canonical) for a KV pool's paged-cache groups. + + Publish groups iff the scheduler ext is flat-built (a radix ext never + delivers flat tables — capture would bind dead buffers). Speculative + decoding is supported: verify writes per-group [bs*N] locations and the + drafter consumes the full-attention group's table (mirrored into + req_to_page each step). Publication is THE upstream signal every flat + consumer keys off. + + Args: + layer_types: Per-layer paged-cache labels (empty -> single + full-history group). + sliding_window_tokens / page_size: Forwarded to + group_specs_from_layer_types. + max_live_requests / max_scheduled_tokens / max_total_tokens / + max_context_len: Sizing inputs for + compute_paged_cache_group_page_counts. + + Returns: + (specs, page_counts) when publishing, None on a radix ext. + """ + if not scheduler_ext_flat_kvcache(): + return None + specs = group_specs_from_layer_types( + layer_types=tuple(layer_types) or (FULL_ATTENTION,), + sliding_window_tokens=sliding_window_tokens, + page_size=page_size, + ) + counts = compute_paged_cache_group_page_counts( + specs, + max_live_requests=max_live_requests, + max_scheduled_tokens=max(0, int(max_scheduled_tokens)), + max_total_tokens=max_total_tokens, + max_context_len=max_context_len, + ) + return specs, counts + + +def compute_max_logical_pages_for_capture( + spec: PagedCacheGroupSpec, + *, + max_context_len: int, + max_tokens_per_req: int = 1, + overlap_schedule_depth: int = 0, +) -> int: + """Return CUDA Graph block-table width for one paged-cache group. + + Decode admission reserves the current verify span plus one span for each + overlapped schedule. Include that complete reservation horizon here: a + request close to the model context limit can still expose the reserved + pages in its scheduler block-table row before the accepted tokens are + truncated by the request-length limit. + + Args: + spec: Paged-cache group layout and retention policy. + max_context_len: Maximum accepted raw-token context length. + max_tokens_per_req: Runtime decode/verify width. + overlap_schedule_depth: Number of additionally in-flight decode steps. + + Returns: + Required block-table columns for one request. + """ + # Local import: keeps this module torch-free at import time. + from tokenspeed.runtime.utils.common import ceil_div + + if max_context_len < 0: + raise ValueError(f"max_context_len must be >= 0, got {max_context_len}") + if max_tokens_per_req <= 0: + raise ValueError(f"max_tokens_per_req must be > 0, got {max_tokens_per_req}") + if overlap_schedule_depth not in (0, 1): + raise ValueError( + f"overlap_schedule_depth must be 0 or 1, got {overlap_schedule_depth}" + ) + raw_per_page = spec.rows_per_page * spec.entry_stride_tokens + if raw_per_page <= 0: + raise ValueError( + f"PagedCacheGroupSpec {spec.group_id}: rows_per_page * " + "entry_stride_tokens must be > 0" + ) + reservation_horizon = (overlap_schedule_depth + 1) * max_tokens_per_req + if spec.retention == "sliding_window": + window = spec.sliding_window_tokens + if window is None or window <= 0: + raise ValueError( + f"PagedCacheGroupSpec {spec.group_id}: sliding group missing " + "positive sliding_window_tokens" + ) + # Capture uses a conservative metadata bound; it does not change the + # per-token attention history counted as window - 1 above. + retention_bound = min(window, max_context_len) + live_tokens = retention_bound + reservation_horizon + return ceil_div(live_tokens, raw_per_page) + 1 + if spec.retention == "full_history": + live_tokens = max_context_len + reservation_horizon + return ceil_div(live_tokens, raw_per_page) + raise ValueError( + f"PagedCacheGroupSpec {spec.group_id}: unsupported retention " + f"{spec.retention!r}" + ) + + +__all__ = [ + "FULL_ATTENTION", + "LINEAR_ATTENTION", + "PagedCacheGroupSpec", + "Retention", + "STATE_LAYER_TYPES", + "compute_max_logical_pages_for_capture", + "compute_paged_cache_group_page_counts", + "group_specs_from_layer_types", + "hybrid_slab_group_size", + "layer_group_ids", + "publish_paged_cache_groups", + "scheduler_ext_flat_kvcache", + "validate_flat_scheduler_config", +] diff --git a/python/tokenspeed/runtime/configs/qwen2_config.py b/python/tokenspeed/runtime/configs/qwen2_config.py new file mode 100644 index 0000000..279b4f6 --- /dev/null +++ b/python/tokenspeed/runtime/configs/qwen2_config.py @@ -0,0 +1,109 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Qwen2 model configuration definitions.""" + +import logging + +from transformers.configuration_utils import PretrainedConfig +from transformers.modeling_rope_utils import rope_config_validation + +logger = logging.getLogger(__name__) + + +class Qwen2Config(PretrainedConfig): + r""" + Configuration class for Qwen2 / Qwen2.5 dense LMs (e.g. ``Qwen/Qwen2-7B``, + ``Qwen/Qwen2.5-7B``). Mirrors :class:`Qwen3Config` but without per-head + Q/K RMSNorm and with Qwen2's fixed attention bias convention + (``q_proj``/``k_proj``/``v_proj`` carry biases, ``o_proj`` does not). + """ + + model_type = "qwen2" + keys_to_ignore_at_inference = ["past_key_values"] + + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + + def __init__( + self, + vocab_size=151936, + hidden_size=4096, + intermediate_size=22016, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=32, + 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, + use_sliding_window=False, + sliding_window=4096, + max_window_layers=28, + attention_dropout=0.0, + **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.use_sliding_window = use_sliding_window + self.sliding_window = sliding_window if self.use_sliding_window else None + self.max_window_layers = max_window_layers + + 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.attention_dropout = attention_dropout + if self.rope_scaling is not None and "type" in self.rope_scaling: + self.rope_scaling["rope_type"] = self.rope_scaling["type"] + rope_config_validation(self) + + super().__init__( + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + +__all__ = ["Qwen2Config"] diff --git a/python/tokenspeed/runtime/configs/qwen3_5_config.py b/python/tokenspeed/runtime/configs/qwen3_5_config.py new file mode 100644 index 0000000..ee2328e --- /dev/null +++ b/python/tokenspeed/runtime/configs/qwen3_5_config.py @@ -0,0 +1,168 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Qwen 3.5 configuration wrappers.""" + +from transformers import PretrainedConfig + +from tokenspeed.runtime.configs.qwen3_5_text_base_config import Qwen3_5BaseTextConfig +from tokenspeed.runtime.configs.qwen3_vision_config import Qwen3VLVisionConfig + +_MROPE_EXTENSION_KEYS = frozenset({"mrope_section", "mrope_interleaved"}) + + +def _to_transformers_rope_parameters(rope_config): + if not isinstance(rope_config, dict): + return rope_config + return { + key: value + for key, value in rope_config.items() + if key not in _MROPE_EXTENSION_KEYS + } + + +class Qwen3_5VisionConfig(Qwen3VLVisionConfig): + model_type = "qwen3_5" + base_config_key = "vision_config" + + +class Qwen3_5TextConfig(Qwen3_5BaseTextConfig): + model_type = "qwen3_5_text" + base_config_key = "text_config" + + def __init__( + self, + **kwargs, + ): + # HF Qwen3.5 checkpoints may provide RoPE settings under rope_parameters. + # Normalize it before parent init so downstream code sees the expected values. + full_rope_parameters = kwargs.get("rope_parameters") + if full_rope_parameters is not None: + kwargs["rope_parameters"] = _to_transformers_rope_parameters( + full_rope_parameters + ) + + super().__init__(**kwargs) + self._tokenspeed_rope_parameters = full_rope_parameters or self.rope_parameters + + +class Qwen3_5Config(PretrainedConfig): + model_type = "qwen3_5" + sub_configs = { + "vision_config": Qwen3_5VisionConfig, + "text_config": Qwen3_5TextConfig, + } + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + text_config=None, + vision_config=None, + image_token_id=151655, + video_token_id=151656, + vision_start_token_id=151652, + vision_end_token_id=151653, + tie_word_embeddings=False, + **kwargs, + ): + self.vision_config = self._ensure_vision_config(vision_config) + self.text_config = self._ensure_text_config(text_config) + + self.image_token_id = image_token_id + self.video_token_id = video_token_id + self.vision_start_token_id = vision_start_token_id + self.vision_end_token_id = vision_end_token_id + super().__init__(**kwargs, tie_word_embeddings=tie_word_embeddings) + + def _ensure_text_config(self, text_config): + """Convert text_config to the proper config class if it's a dict.""" + text_cls = self.sub_configs["text_config"] + if isinstance(text_config, dict): + return text_cls(**text_config) + if isinstance(text_config, Qwen3_5Config): + nested_text_config = text_config.__dict__.get("text_config") + if nested_text_config is not None and nested_text_config is not text_config: + return self._ensure_text_config(nested_text_config) + if text_config is None: + return text_cls() + return text_config + + def _ensure_vision_config(self, vision_config): + """Convert vision_config to the proper config class if it's a dict.""" + vision_cls = self.sub_configs["vision_config"] + if isinstance(vision_config, dict): + return vision_cls(**vision_config) + if vision_config is None: + return vision_cls() + return vision_config + + def __setattr__(self, name, value): + # from_pretrained re-assigns text_config as a raw dict after __init__; + # intercept and convert it back to the proper config class. + if name == "text_config" and isinstance(value, dict): + value = self._ensure_text_config(value) + elif name == "vision_config" and isinstance(value, dict): + value = self._ensure_vision_config(value) + super().__setattr__(name, value) + + def __getattr__(self, name): + """Forward attribute access to text_config for inference-only usage.""" + if name.startswith("_") or name in {"text_config", "vision_config"}: + raise AttributeError(f"'{type(self).__name__}' has no attribute '{name}'") + + text_config = self.__dict__.get("text_config") + if isinstance(text_config, dict): + text_config = self._ensure_text_config(text_config) + self.__dict__["text_config"] = text_config + if text_config is None or text_config is self: + raise AttributeError(f"'{type(self).__name__}' has no attribute '{name}'") + + try: + return getattr(text_config, name) + except AttributeError as exc: + raise AttributeError( + f"'{type(self).__name__}' has no attribute '{name}'" + ) from exc + + +class Qwen3_5MoeVisionConfig(Qwen3_5VisionConfig): + model_type = "qwen3_5_moe" + + +class Qwen3_5MoeTextConfig(Qwen3_5TextConfig): + model_type = "qwen3_5_moe_text" + + def __init__(self, **kwargs): + # Explicit __init__ prevents transformers from auto-generating one + # that skips Qwen3_5TextConfig.__init__ (rope_parameters normalization). + super().__init__(**kwargs) + + +class Qwen3_5MoeConfig(Qwen3_5Config): + model_type = "qwen3_5_moe" + sub_configs = { + "vision_config": Qwen3_5MoeVisionConfig, + "text_config": Qwen3_5MoeTextConfig, + } + + def __init__(self, **kwargs): + # Explicit __init__ prevents transformers from auto-generating one + # that skips Qwen3_5Config.__init__ (text/vision config setup). + super().__init__(**kwargs) diff --git a/python/tokenspeed/runtime/configs/qwen3_5_text_base_config.py b/python/tokenspeed/runtime/configs/qwen3_5_text_base_config.py new file mode 100755 index 0000000..f351259 --- /dev/null +++ b/python/tokenspeed/runtime/configs/qwen3_5_text_base_config.py @@ -0,0 +1,359 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright 2025 The Qwen team, Alibaba Group +# SPDX-FileCopyrightText: Copyright 2025 HuggingFace Inc. team +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Qwen3.5 text model configuration definitions.""" + +import enum + +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 tokenspeed.runtime.configs.paged_cache_spec import FULL_ATTENTION +from tokenspeed.runtime.distributed.utils import divide +from tokenspeed.runtime.utils.env import envs + +logger = logging.get_logger(__name__) + + +# HybridLayerType +class HybridLayerType(enum.Enum): + full_attention = "attention" + swa_attention = "swa_attention" + linear_attention = "linear_attention" + mamba2 = "mamba" + + +class Qwen3_5BaseTextConfig(PretrainedConfig): + r""" + Shared text configuration base used by Qwen3.5 dense and MoE configs. + + 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_parameters (`Dict`, *optional*): + Dictionary containing the scaling configuration for the RoPE embeddings. 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 dense MLP rather than sparse MoE blocks. + 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 Qwen3_5TextModel, Qwen3_5BaseTextConfig + + >>> # Initializing a Qwen3.5 text configuration + >>> configuration = Qwen3_5BaseTextConfig() + + >>> # Initializing a model from the Qwen3.5 text-80B-A3B style configuration + >>> model = Qwen3_5TextModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "qwen3_5_text_base" + 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_parameters=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_parameters = rope_parameters + 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 layer_index in range(self.num_hidden_layers): + if (layer_index + 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 layer_types(self): + """Per-layer paged-cache labels: "full_attention" / "linear_attention". + + Same interleaving as ``layers_block_type``, translated to the label + vocabulary of ``paged_cache_spec`` (``HybridLayerType.full_attention`` + serializes as the checkpoint's "attention", which the KV-cache layer + has no retention entry for). A property rather than an ``__init__`` + attribute because NextN drafts overwrite ``full_attention_interval`` + after construction (models/qwen3_5_nextn.py). + """ + return [ + ( + FULL_ATTENTION + if layer_type == HybridLayerType.full_attention.value + else layer_type + ) + for layer_type in self.layers_block_type + ] + + @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 mamba2_cache_params(self): + # Imported lazily to avoid config/env import cycles during module initialization. + from tokenspeed.runtime.utils.env import global_server_args_dict + + self.mapping = global_server_args_dict["mapping"] + attn_tp_size = self.mapping.attn.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, attn_tp_size), + self.linear_conv_kernel_dim - 1, + ) + + temporal_state_shape = ( + divide(self.linear_num_value_heads, attn_tp_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[envs.TOKENSPEED_MAMBA_SSM_DTYPE.get()] + 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.mamba2_cache_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/tokenspeed/runtime/configs/qwen3_asr_config.py b/python/tokenspeed/runtime/configs/qwen3_asr_config.py new file mode 100644 index 0000000..52202bc --- /dev/null +++ b/python/tokenspeed/runtime/configs/qwen3_asr_config.py @@ -0,0 +1,204 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Qwen3-ASR nested configuration definitions. + +Transformers 5.12 contains the shared Qwen3-Omni audio architecture but does +not yet register ``model_type: qwen3_asr``. These small wrappers preserve the +official checkpoint's ``thinker_config`` shape while reusing TokenSpeed's +Qwen3 text configuration. +""" + +from __future__ import annotations + +from typing import Any + +from transformers import PretrainedConfig + +from tokenspeed.runtime.configs.qwen3_config import Qwen3Config + + +class Qwen3ASRAudioEncoderConfig(PretrainedConfig): + """Qwen3 audio tower configuration shipped by Qwen3-ASR.""" + + model_type = "qwen3_asr_audio_encoder" + base_config_key = "audio_config" + attribute_map = { + "num_hidden_layers": "encoder_layers", + "hidden_size": "d_model", + "num_attention_heads": "encoder_attention_heads", + "intermediate_size": "encoder_ffn_dim", + } + + def __init__( + self, + *, + num_mel_bins: int = 128, + encoder_layers: int = 24, + encoder_attention_heads: int = 16, + encoder_ffn_dim: int = 4096, + d_model: int = 1024, + dropout: float = 0.0, + attention_dropout: float = 0.0, + activation_function: str = "gelu", + activation_dropout: float = 0.0, + scale_embedding: bool = False, + initializer_range: float = 0.02, + max_source_positions: int = 1500, + n_window: int = 50, + n_window_infer: int = 800, + conv_chunksize: int = 500, + downsample_hidden_size: int = 480, + output_dim: int = 2048, + **kwargs: Any, + ) -> None: + self.num_mel_bins = num_mel_bins + self.encoder_layers = encoder_layers + self.encoder_attention_heads = encoder_attention_heads + self.encoder_ffn_dim = encoder_ffn_dim + self.d_model = d_model + self.dropout = dropout + self.attention_dropout = attention_dropout + self.activation_function = activation_function + self.activation_dropout = activation_dropout + self.scale_embedding = scale_embedding + self.initializer_range = initializer_range + self.max_source_positions = max_source_positions + self.n_window = n_window + self.n_window_infer = n_window_infer + self.conv_chunksize = conv_chunksize + self.downsample_hidden_size = downsample_hidden_size + self.output_dim = output_dim + super().__init__(**kwargs) + + +class Qwen3ASRThinkerConfig(PretrainedConfig): + """Audio/text thinker configuration nested below the outer ASR config.""" + + model_type = "qwen3_asr_thinker" + base_config_key = "thinker_config" + keys_to_ignore_at_inference = ["past_key_values"] + sub_configs = { + "audio_config": Qwen3ASRAudioEncoderConfig, + "text_config": Qwen3Config, + } + + def __init__( + self, + *, + audio_config: dict[str, Any] | Qwen3ASRAudioEncoderConfig | None = None, + text_config: dict[str, Any] | Qwen3Config | None = None, + audio_start_token_id: int = 151669, + audio_end_token_id: int = 151670, + audio_token_id: int = 151676, + initializer_range: float = 0.02, + **kwargs: Any, + ) -> None: + self.audio_config = self._ensure_audio_config(audio_config) + self.text_config = self._ensure_text_config(text_config) + self.audio_start_token_id = audio_start_token_id + self.audio_end_token_id = audio_end_token_id + self.audio_token_id = audio_token_id + self.initializer_range = initializer_range + super().__init__(**kwargs) + + @staticmethod + def _ensure_audio_config( + config: dict[str, Any] | Qwen3ASRAudioEncoderConfig | None, + ) -> Qwen3ASRAudioEncoderConfig: + if isinstance(config, Qwen3ASRAudioEncoderConfig): + return config + return Qwen3ASRAudioEncoderConfig(**(config or {})) + + @staticmethod + def _ensure_text_config( + config: dict[str, Any] | Qwen3Config | None, + ) -> Qwen3Config: + if isinstance(config, Qwen3Config): + return config + return Qwen3Config(**(config or {})) + + def __setattr__(self, name: str, value: Any) -> None: + if name == "audio_config" and isinstance(value, dict): + value = self._ensure_audio_config(value) + elif name == "text_config" and isinstance(value, dict): + value = self._ensure_text_config(value) + super().__setattr__(name, value) + + def get_text_config(self, *args: Any, **kwargs: Any) -> Qwen3Config: + return self.text_config + + +class Qwen3ASRConfig(PretrainedConfig): + """Outer config matching ``Qwen/Qwen3-ASR-*`` config.json files.""" + + model_type = "qwen3_asr" + keys_to_ignore_at_inference = ["past_key_values"] + sub_configs = {"thinker_config": Qwen3ASRThinkerConfig} + + def __init__( + self, + *, + thinker_config: dict[str, Any] | Qwen3ASRThinkerConfig | None = None, + support_languages: list[str] | None = None, + **kwargs: Any, + ) -> None: + self.thinker_config = self._ensure_thinker_config(thinker_config) + self.support_languages = list(support_languages or []) + super().__init__(**kwargs) + + @staticmethod + def _ensure_thinker_config( + config: dict[str, Any] | Qwen3ASRThinkerConfig | None, + ) -> Qwen3ASRThinkerConfig: + if isinstance(config, Qwen3ASRThinkerConfig): + return config + return Qwen3ASRThinkerConfig(**(config or {})) + + def __setattr__(self, name: str, value: Any) -> None: + if name == "thinker_config" and isinstance(value, dict): + value = self._ensure_thinker_config(value) + super().__setattr__(name, value) + + def get_text_config(self, *args: Any, **kwargs: Any) -> Qwen3Config: + return self.thinker_config.text_config + + @property + def audio_config(self) -> Qwen3ASRAudioEncoderConfig: + return self.thinker_config.audio_config + + @property + def audio_start_token_id(self) -> int: + return self.thinker_config.audio_start_token_id + + @property + def audio_end_token_id(self) -> int: + return self.thinker_config.audio_end_token_id + + @property + def audio_token_id(self) -> int: + return self.thinker_config.audio_token_id + + +__all__ = [ + "Qwen3ASRAudioEncoderConfig", + "Qwen3ASRConfig", + "Qwen3ASRThinkerConfig", +] diff --git a/python/tokenspeed/runtime/configs/qwen3_config.py b/python/tokenspeed/runtime/configs/qwen3_config.py new file mode 100755 index 0000000..a940284 --- /dev/null +++ b/python/tokenspeed/runtime/configs/qwen3_config.py @@ -0,0 +1,246 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Qwen3 model configuration definitions.""" + +import logging + +from transformers.configuration_utils import PretrainedConfig +from transformers.modeling_rope_utils import rope_config_validation + +ALLOWED_LAYER_TYPES = ( + "full_attention", + "sliding_attention", + "chunked_attention", +) + + +def layer_type_validation(layer_types: list[str]): + """Check that each entry in `layer_types` are allowed.""" + if not all(layer_type in ALLOWED_LAYER_TYPES for layer_type in layer_types): + raise ValueError(f"The `layer_types` entries must be in {ALLOWED_LAYER_TYPES}") + + +logger = logging.getLogger(__name__) + + +class Qwen3Config(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a [`Qwen3Model`]. It is used to instantiate a + Qwen3 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-8B [Qwen/Qwen3-8B](https://huggingface.co/Qwen/Qwen3-8B). + + 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 Qwen3 model. Defines the number of different tokens that can be represented by the + `inputs_ids` passed when calling [`Qwen3Model`] + hidden_size (`int`, *optional*, defaults to 4096): + Dimension of the hidden representations. + intermediate_size (`int`, *optional*, defaults to 22016): + 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*, defaults to 32): + 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`. + head_dim (`int`, *optional*, defaults to 128): + The attention head dimension. + 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 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. 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 + attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`): + Whether to use a bias in the query, key, value and output projection layers during self-attention. + use_sliding_window (`bool`, *optional*, defaults to `False`): + Whether to use sliding window attention. + sliding_window (`int`, *optional*, defaults to 4096): + Sliding window attention (SWA) window size. If not specified, will default to `4096`. + max_window_layers (`int`, *optional*, defaults to 28): + The number of layers that use SWA (Sliding Window Attention). The bottom layers use SWA while the top use full attention. + layer_types (`list`, *optional*): + Attention pattern for each layer. + attention_dropout (`float`, *optional*, defaults to 0.0): + The dropout ratio for the attention probabilities. + + ```python + >>> from transformers import Qwen3Model, Qwen3Config + + >>> # Initializing a Qwen3 style configuration + >>> configuration = Qwen3Config() + + >>> # Initializing a model from the Qwen3-8B style configuration + >>> model = Qwen3Model(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "qwen3" + keys_to_ignore_at_inference = ["past_key_values"] + + # Default tensor parallel plan for base model `Qwen3` + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise", + "layers.*.self_attn.k_proj": "colwise", + "layers.*.self_attn.v_proj": "colwise", + "layers.*.self_attn.o_proj": "rowwise", + "layers.*.mlp.gate_proj": "colwise", + "layers.*.mlp.up_proj": "colwise", + "layers.*.mlp.down_proj": "rowwise", + } + base_model_pp_plan = { + "embed_tokens": (["input_ids"], ["inputs_embeds"]), + "layers": (["hidden_states", "attention_mask"], ["hidden_states"]), + "norm": (["hidden_states"], ["hidden_states"]), + } + + def __init__( + self, + vocab_size=151936, + hidden_size=4096, + intermediate_size=22016, + num_hidden_layers=32, + num_attention_heads=32, + num_key_value_heads=32, + head_dim=128, + 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, + attention_bias=False, + use_sliding_window=False, + sliding_window=4096, + max_window_layers=28, + layer_types=None, + attention_dropout=0.0, + **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.use_sliding_window = use_sliding_window + self.sliding_window = sliding_window if self.use_sliding_window else None + self.max_window_layers = max_window_layers + + # 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.head_dim = head_dim + 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.attention_bias = attention_bias + self.attention_dropout = attention_dropout + # Validate the correctness of rotary position embeddings parameters + # BC: if there is a 'type' field, move it to 'rope_type'. + if self.rope_scaling is not None and "type" in self.rope_scaling: + self.rope_scaling["rope_type"] = self.rope_scaling["type"] + rope_config_validation(self) + + self.layer_types = layer_types + if self.layer_types is None: + self.layer_types = [ + ( + "sliding_attention" + if self.sliding_window is not None and i >= self.max_window_layers + else "full_attention" + ) + for i in range(self.num_hidden_layers) + ] + layer_type_validation(self.layer_types) + + super().__init__( + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + +__all__ = ["Qwen3Config"] diff --git a/python/tokenspeed/runtime/configs/qwen3_moe_config.py b/python/tokenspeed/runtime/configs/qwen3_moe_config.py new file mode 100644 index 0000000..1b4e9d4 --- /dev/null +++ b/python/tokenspeed/runtime/configs/qwen3_moe_config.py @@ -0,0 +1,56 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Qwen3 MoE model configuration definitions.""" + +from tokenspeed.runtime.configs.qwen3_config import Qwen3Config + + +class Qwen3MoeConfig(Qwen3Config): + """Configuration for Qwen3 MoE causal LMs such as Qwen3-30B-A3B.""" + + model_type = "qwen3_moe" + + def __init__( + self, + decoder_sparse_step=1, + moe_intermediate_size=768, + shared_expert_intermediate_size=0, + num_experts_per_tok=8, + num_experts=128, + norm_topk_prob=True, + output_router_logits=False, + router_aux_loss_coef=0.001, + mlp_only_layers=None, + **kwargs, + ): + super().__init__(**kwargs) + 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 = [] if mlp_only_layers is None else mlp_only_layers + + +__all__ = ["Qwen3MoeConfig"] diff --git a/python/tokenspeed/runtime/configs/qwen3_vision_config.py b/python/tokenspeed/runtime/configs/qwen3_vision_config.py new file mode 100644 index 0000000..b50f58a --- /dev/null +++ b/python/tokenspeed/runtime/configs/qwen3_vision_config.py @@ -0,0 +1,59 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from transformers import PretrainedConfig + + +class Qwen3VLVisionConfig(PretrainedConfig): + model_type = "qwen3_vl" + base_config_key = "vision_config" + + def __init__( + self, + depth=27, + hidden_size=1152, + hidden_act="gelu_pytorch_tanh", + intermediate_size=4304, + num_heads=16, + in_channels=3, + patch_size=16, + spatial_merge_size=2, + temporal_patch_size=2, + out_hidden_size=3584, + num_position_embeddings=2304, + deepstack_visual_indexes=[8, 16, 24], + initializer_range=0.02, + **kwargs, + ): + super().__init__(**kwargs) + + self.depth = depth + self.hidden_size = hidden_size + self.hidden_act = hidden_act + self.intermediate_size = intermediate_size + self.num_heads = num_heads + self.in_channels = in_channels + self.patch_size = patch_size + self.spatial_merge_size = spatial_merge_size + self.temporal_patch_size = temporal_patch_size + self.out_hidden_size = out_hidden_size + self.num_position_embeddings = num_position_embeddings + self.initializer_range = initializer_range + self.deepstack_visual_indexes = deepstack_visual_indexes diff --git a/python/tokenspeed/runtime/configs/utils.py b/python/tokenspeed/runtime/configs/utils.py new file mode 100755 index 0000000..64658db --- /dev/null +++ b/python/tokenspeed/runtime/configs/utils.py @@ -0,0 +1,809 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Shared configuration utilities.""" + +import math + +from transformers.configuration_utils import PretrainedConfig +from transformers.utils import is_torch_available, logging + +logger = logging.get_logger(__name__) + + +if is_torch_available(): + import torch + + +def get_rope_theta(config, default: float = 10000.0) -> float: + """Return rope_theta from config, including the transformers 5.x fallback.""" + theta = getattr(config, "rope_theta", None) + if theta is not None: + return theta + rope_scaling = getattr(config, "rope_scaling", None) + if isinstance(rope_scaling, dict): + return rope_scaling.get("rope_theta", default) + return default + + +def get_rope_parameters(config): + """Return TokenSpeed's full RoPE config, including private extensions.""" + return ( + getattr(config, "_tokenspeed_rope_parameters", None) + or getattr(config, "rope_parameters", None) + or {} + ) + + +def _compute_default_rope_parameters( + config: PretrainedConfig | None = None, + device: "torch.device | None" = None, + seq_len: int | None = None, + **rope_kwargs, +) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies according to the original RoPE implementation + Args: + config ([`~transformers.PretrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. Unused for this type of RoPE. + rope_kwargs (`Dict`, *optional*): + BC compatibility with the previous RoPE class instantiation, will be removed in v4.45. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). + """ + if config is not None and rope_kwargs: + raise ValueError( + "Unexpected arguments: `**rope_kwargs` and `config` are mutually exclusive in " + f"`_compute_default_rope_parameters`, got `rope_kwargs`={rope_kwargs} and `config`={config}" + ) + if rope_kwargs: + base = rope_kwargs["base"] + dim = rope_kwargs["dim"] + elif config is not None: + base = config.rope_theta + partial_rotary_factor = ( + config.partial_rotary_factor + if hasattr(config, "partial_rotary_factor") + else 1.0 + ) + head_dim = getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads + ) + dim = int(head_dim * partial_rotary_factor) + + attention_factor = 1.0 # Unused in this type of RoPE + + # Compute the inverse frequencies + inv_freq = 1.0 / ( + base + ** ( + torch.arange(0, dim, 2, dtype=torch.int64).to( + device=device, dtype=torch.float + ) + / dim + ) + ) + return inv_freq, attention_factor + + +def _compute_linear_scaling_rope_parameters( + config: PretrainedConfig | None = None, + device: "torch.device | None" = None, + seq_len: int | None = None, + **rope_kwargs, +) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies with linear scaling. Credits to the Reddit user /u/kaiokendev + Args: + config ([`~transformers.PretrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. Unused for this type of RoPE. + rope_kwargs (`Dict`, *optional*): + BC compatibility with the previous RoPE class instantiation, will be removed in v4.45. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). + """ + if config is not None and rope_kwargs: + raise ValueError( + "Unexpected arguments: `**rope_kwargs` and `config` are mutually exclusive in " + f"`_compute_linear_scaling_rope_parameters`, got `rope_kwargs`={rope_kwargs} and `config`={config}" + ) + if rope_kwargs: + factor = rope_kwargs["factor"] + elif config is not None: + factor = config.rope_scaling["factor"] + + # Gets the default RoPE parameters + inv_freq, attention_factor = _compute_default_rope_parameters( + config, device, seq_len, **rope_kwargs + ) + + # Then applies linear scaling to the frequencies. + # originally, scaling was applied to the position_ids. However, we get `embs = inv_freq @ position_ids`, so + # applying scaling to the inverse frequencies is equivalent. + inv_freq /= factor + return inv_freq, attention_factor + + +def _compute_dynamic_ntk_parameters( + config: PretrainedConfig | None = None, + device: "torch.device | None" = None, + seq_len: int | None = None, + **rope_kwargs, +) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies with NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla + Args: + config ([`~transformers.PretrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length, used to update the dynamic RoPE at inference time. + rope_kwargs (`Dict`, *optional*): + BC compatibility with the previous RoPE class instantiation, will be removed in v4.45. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). + """ + if config is not None and rope_kwargs: + raise ValueError( + "Unexpected arguments: `**rope_kwargs` and `config` are mutually exclusive in " + f"`_compute_dynamic_ntk_parameters`, got `rope_kwargs`={rope_kwargs} and `config`={config}" + ) + if rope_kwargs: + base = rope_kwargs["base"] + dim = rope_kwargs["dim"] + max_position_embeddings = rope_kwargs["max_position_embeddings"] + factor = rope_kwargs["factor"] + elif config is not None: + base = config.rope_theta + partial_rotary_factor = ( + config.partial_rotary_factor + if hasattr(config, "partial_rotary_factor") + else 1.0 + ) + head_dim = getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads + ) + dim = int(head_dim * partial_rotary_factor) + max_position_embeddings = config.max_position_embeddings + factor = config.rope_scaling["factor"] + + attention_factor = 1.0 # Unused in this type of RoPE + + # seq_len: default to max_position_embeddings, e.g. at init time + seq_len = ( + seq_len + if seq_len is not None and seq_len > max_position_embeddings + else max_position_embeddings + ) + + # Compute the inverse frequencies + base = base * ((factor * seq_len / max_position_embeddings) - (factor - 1)) ** ( + dim / (dim - 2) + ) + inv_freq = 1.0 / ( + base + ** ( + torch.arange(0, dim, 2, dtype=torch.int64).to( + device=device, dtype=torch.float + ) + / dim + ) + ) + return inv_freq, attention_factor + + +def _compute_yarn_parameters( + config: PretrainedConfig, + device: "torch.device", + seq_len: int | None = None, + **rope_kwargs, +) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies with NTK scaling. Please refer to the + [original paper](https://arxiv.org/abs/2309.00071) + Args: + config ([`~transformers.PretrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. Unused for this type of RoPE. + rope_kwargs (`Dict`, *optional*): + BC compatibility with the previous RoPE class instantiation, will be removed in v4.45. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin. + """ + # No need to keep BC with yarn, unreleased when this new pattern was created. + if rope_kwargs: + raise ValueError( + f"Unexpected arguments: `**rope_kwargs` should be unset in `_compute_yarn_parameters`, got {rope_kwargs}" + ) + + base = config.rope_theta + partial_rotary_factor = ( + config.partial_rotary_factor + if hasattr(config, "partial_rotary_factor") + else 1.0 + ) + head_dim = getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads + ) + dim = int(head_dim * partial_rotary_factor) + factor = config.rope_scaling["factor"] + attention_factor = config.rope_scaling.get("attention_factor") + mscale = config.rope_scaling.get("mscale") + mscale_all_dim = config.rope_scaling.get("mscale_all_dim") + + # DeekSeek-V3 (and potentially other models) modify `max_position_embeddings` and have a + # `original_max_position_embeddings` field containing the pretrained value. They use the ratio between these two + # values to compute the default attention scaling factor, instead of using `factor`. + if "original_max_position_embeddings" in config.rope_scaling: + original_max_position_embeddings = config.rope_scaling[ + "original_max_position_embeddings" + ] + factor = config.max_position_embeddings / original_max_position_embeddings + else: + original_max_position_embeddings = config.max_position_embeddings + + def get_mscale(scale, mscale=1): + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + # Sets the attention factor as suggested in the paper + if attention_factor is None: + if mscale and mscale_all_dim: + attention_factor = float( + get_mscale(factor, mscale) / get_mscale(factor, mscale_all_dim) + ) + else: + attention_factor = get_mscale(factor) + + # Optional config options + # beta_fast/beta_slow: as suggested in the paper, default to 32/1 (correspondingly) + beta_fast = config.rope_scaling.get("beta_fast") or 32 + beta_slow = config.rope_scaling.get("beta_slow") or 1 + + # Compute the inverse frequencies + def find_correction_dim(num_rotations, dim, base, max_position_embeddings): + """Inverse dimension formula to find the dimension based on the number of rotations""" + return ( + dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi)) + ) / (2 * math.log(base)) + + def find_correction_range(low_rot, high_rot, dim, base, max_position_embeddings): + """Find dimension range bounds based on rotations""" + low = math.floor( + find_correction_dim(low_rot, dim, base, max_position_embeddings) + ) + high = math.ceil( + find_correction_dim(high_rot, dim, base, max_position_embeddings) + ) + return max(low, 0), min(high, dim - 1) + + def linear_ramp_factor(min, max, dim): + if min == max: + max += 0.001 # Prevent singularity + + linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min) + ramp_func = torch.clamp(linear_func, 0, 1) + return ramp_func + + # Note on variable naming: "interpolation" comes from the original technique, where we interpolate the position IDs + # to expand the possible context length. In other words, interpolation = apply scaling factor. + pos_freqs = base ** ( + torch.arange(0, dim, 2).to(device=device, dtype=torch.float) / dim + ) + inv_freq_extrapolation = 1.0 / pos_freqs + inv_freq_interpolation = 1.0 / (factor * pos_freqs) + + low, high = find_correction_range( + beta_fast, beta_slow, dim, base, original_max_position_embeddings + ) + + # Get n-dimensional rotational scaling corrected for extrapolation + inv_freq_extrapolation_factor = 1 - linear_ramp_factor(low, high, dim // 2).to( + device=device, dtype=torch.float + ) + inv_freq = ( + inv_freq_interpolation * (1 - inv_freq_extrapolation_factor) + + inv_freq_extrapolation * inv_freq_extrapolation_factor + ) + return inv_freq, attention_factor + + +def _compute_longrope_parameters( + config: PretrainedConfig, + device: "torch.device", + seq_len: int | None = None, + **rope_kwargs, +) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies with LongRoPE scaling. Please refer to the + [original implementation](https://github.com/microsoft/LongRoPE) + Args: + config ([`~transformers.PretrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. + rope_kwargs (`Dict`, *optional*): + BC compatibility with the previous RoPE class instantiation, will be removed in v4.45. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin. + """ + # No need to keep BC with longrope, unreleased when this new pattern was created. + if rope_kwargs: + raise ValueError( + "Unexpected arguments: `**rope_kwargs` should be unset in `_compute_longrope_parameters`, got " + f"{rope_kwargs}" + ) + + base = config.rope_theta + partial_rotary_factor = ( + config.partial_rotary_factor + if hasattr(config, "partial_rotary_factor") + else 1.0 + ) + head_dim = getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads + ) + dim = int(head_dim * partial_rotary_factor) + long_factor = config.rope_scaling["long_factor"] + short_factor = config.rope_scaling["short_factor"] + factor = config.rope_scaling.get("factor") + attention_factor = config.rope_scaling.get("attention_factor") + + # Phi3 (and potentially other models) modify `max_position_embeddings` and have a + # `original_max_position_embeddings` field containing the pretrained value. They use the ratio between these two + # values to compute the default attention scaling factor, instead of using `factor`. + if hasattr(config, "original_max_position_embeddings"): + original_max_position_embeddings = config.original_max_position_embeddings + factor = ( + config.max_position_embeddings / config.original_max_position_embeddings + ) + else: + original_max_position_embeddings = config.max_position_embeddings + + # Sets the attention factor as suggested in the paper + if attention_factor is None: + if factor <= 1.0: + attention_factor = 1.0 + else: + attention_factor = math.sqrt( + 1 + math.log(factor) / math.log(original_max_position_embeddings) + ) + + # Compute the inverse frequencies -- scaled based on the target sequence length + if seq_len and seq_len > original_max_position_embeddings: + ext_factors = torch.tensor(long_factor, dtype=torch.float32, device=device) + else: + ext_factors = torch.tensor(short_factor, dtype=torch.float32, device=device) + inv_freq_shape = ( + torch.arange(0, dim, 2, dtype=torch.int64, device=device).float() / dim + ) + inv_freq = 1.0 / (ext_factors * base**inv_freq_shape) + + return inv_freq, attention_factor + + +def _compute_llama3_parameters( + config: PretrainedConfig, + device: "torch.device", + seq_len: int | None = None, + **rope_kwargs, +) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies for llama 3.1. + + Args: + config ([`~transformers.PretrainedConfig`]): + The model configuration. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. Unused for this type of RoPE. + rope_kwargs (`Dict`, *optional*): + BC compatibility with the previous RoPE class instantiation, will be removed in v4.45. + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin. + """ + # Gets the default RoPE parameters + inv_freq, attention_factor = _compute_default_rope_parameters( + config, device, seq_len, **rope_kwargs + ) + + factor = config.rope_scaling["factor"] # `8` in the original implementation + low_freq_factor = config.rope_scaling[ + "low_freq_factor" + ] # `1` in the original implementation + high_freq_factor = config.rope_scaling[ + "high_freq_factor" + ] # `4` in the original implementation + old_context_len = config.rope_scaling[ + "original_max_position_embeddings" + ] # `8192` in the original implementation + + low_freq_wavelen = old_context_len / low_freq_factor + high_freq_wavelen = old_context_len / high_freq_factor + + wavelen = 2 * math.pi / inv_freq + # wavelen < high_freq_wavelen: do nothing + # wavelen > low_freq_wavelen: divide by factor + inv_freq_llama = torch.where( + wavelen > low_freq_wavelen, inv_freq / factor, inv_freq + ) + # otherwise: interpolate between the two, using a smooth factor + smooth_factor = (old_context_len / wavelen - low_freq_factor) / ( + high_freq_factor - low_freq_factor + ) + smoothed_inv_freq = ( + 1 - smooth_factor + ) * inv_freq_llama / factor + smooth_factor * inv_freq_llama + is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen > low_freq_wavelen) + inv_freq_llama = torch.where(is_medium_freq, smoothed_inv_freq, inv_freq_llama) + + return inv_freq_llama, attention_factor + + +# This maps the "rope_type" string field in rope config to the corresponding function to compute the RoPE parameters +# from the model config. You can append new {'rope_type': callable} pairs to this dictionary to enable custom RoPE +# parameterizations, as long as the callable has the same signature. +ROPE_INIT_FUNCTIONS = { + "default": _compute_default_rope_parameters, + "linear": _compute_linear_scaling_rope_parameters, + "dynamic": _compute_dynamic_ntk_parameters, + "yarn": _compute_yarn_parameters, + "longrope": _compute_longrope_parameters, + "llama3": _compute_llama3_parameters, +} + + +def _check_received_keys( + rope_type: str, + received_keys: set, + required_keys: set, + optional_keys: set | None = None, + ignore_keys: set | None = None, +): + """Compare the received keys in `config.rope_scaling` against the expected and optional keys""" + # BC: "rope_type" was originally "type" -- let's check for "rope_type" when "type" is present + if "type" in received_keys: + received_keys -= {"type"} + required_keys.add("rope_type") + + # Some models need to store model-specific keys, and we don't want to throw warning at them + if ignore_keys is not None: + received_keys -= ignore_keys + + missing_keys = required_keys - received_keys + if missing_keys: + raise KeyError( + f"Missing required keys in `rope_scaling` for 'rope_type'='{rope_type}': {missing_keys}" + ) + + if optional_keys is not None: + unused_keys = received_keys - required_keys - optional_keys + else: + unused_keys = received_keys - required_keys + if unused_keys: + logger.warning( + "Unrecognized keys in `rope_scaling` for 'rope_type'='%s': %s", + rope_type, + unused_keys, + ) + + +def _validate_default_rope_parameters( + config: PretrainedConfig, ignore_keys: set | None = None +): + rope_scaling = config.rope_scaling + rope_type = rope_scaling.get( + "rope_type", rope_scaling.get("type", None) + ) # BC: "rope_type" was originally "type" + required_keys = {"rope_type"} + received_keys = set(rope_scaling.keys()) + _check_received_keys( + rope_type, received_keys, required_keys, ignore_keys=ignore_keys + ) + + +def _validate_linear_scaling_rope_parameters( + config: PretrainedConfig, ignore_keys: set | None = None +): + rope_scaling = config.rope_scaling + rope_type = rope_scaling.get( + "rope_type", rope_scaling.get("type", None) + ) # BC: "rope_type" was originally "type" + required_keys = {"rope_type", "factor"} + received_keys = set(rope_scaling.keys()) + _check_received_keys( + rope_type, received_keys, required_keys, ignore_keys=ignore_keys + ) + + factor = rope_scaling["factor"] + if factor is None or not isinstance(factor, float) or factor < 1.0: + logger.warning( + "`rope_scaling`'s factor field must be a float >= 1, got %s", factor + ) + + +def _validate_dynamic_scaling_rope_parameters( + config: PretrainedConfig, ignore_keys: set | None = None +): + rope_scaling = config.rope_scaling + rope_type = rope_scaling.get( + "rope_type", rope_scaling.get("type", None) + ) # BC: "rope_type" was originally "type" + required_keys = {"rope_type", "factor"} + optional_keys = {"original_max_position_embeddings"} + received_keys = set(rope_scaling.keys()) + _check_received_keys( + rope_type, received_keys, required_keys, optional_keys, ignore_keys=ignore_keys + ) + + factor = rope_scaling["factor"] + if factor is None or not isinstance(factor, float) or factor < 1.0: + logger.warning( + "`rope_scaling`'s factor field must be a float >= 1, got %s", factor + ) + + +def _validate_yarn_parameters(config: PretrainedConfig, ignore_keys: set | None = None): + rope_scaling = config.rope_scaling + rope_type = rope_scaling.get( + "rope_type", rope_scaling.get("type", None) + ) # BC: "rope_type" was originally "type" + required_keys = {"rope_type", "factor"} + optional_keys = { + "attention_factor", + "beta_fast", + "beta_slow", + "original_max_position_embeddings", + "mscale", + "mscale_all_dim", + } + received_keys = set(rope_scaling.keys()) + _check_received_keys( + rope_type, received_keys, required_keys, optional_keys, ignore_keys=ignore_keys + ) + + factor = rope_scaling["factor"] + if factor is None or not isinstance(factor, float) or factor < 1.0: + logger.warning( + "`rope_scaling`'s factor field must be a float >= 1, got %s", factor + ) + + attention_factor = rope_scaling.get("attention_factor") + if attention_factor is not None and ( + not isinstance(attention_factor, float) or attention_factor < 0 + ): + logger.warning( + "`rope_scaling`'s attention_factor field must be a float greater than 0, got %s", + attention_factor, + ) + beta_fast = rope_scaling.get("beta_fast") + if beta_fast is not None and not isinstance(beta_fast, float): + logger.warning( + "`rope_scaling`'s beta_fast field must be a float, got %s", beta_fast + ) + beta_slow = rope_scaling.get("beta_slow") + if beta_slow is not None and not isinstance(beta_slow, float): + logger.warning( + "`rope_scaling`'s beta_slow field must be a float, got %s", beta_slow + ) + + if (beta_fast or 32) < (beta_slow or 1): + logger.warning( + "`rope_scaling`'s beta_fast field must be greater than beta_slow, got beta_fast=%s (defaults to 32 if None) and beta_slow=%s (defaults to 1 if None)", + beta_fast, + beta_slow, + ) + + +def _validate_longrope_parameters( + config: PretrainedConfig, ignore_keys: set | None = None +): + rope_scaling = config.rope_scaling + rope_type = rope_scaling.get( + "rope_type", rope_scaling.get("type", None) + ) # BC: "rope_type" was originally "type" + required_keys = {"rope_type", "short_factor", "long_factor"} + optional_keys = {"attention_factor", "factor", "original_max_position_embeddings"} + received_keys = set(rope_scaling.keys()) + _check_received_keys( + rope_type, received_keys, required_keys, optional_keys, ignore_keys=ignore_keys + ) + + partial_rotary_factor = ( + config.partial_rotary_factor + if hasattr(config, "partial_rotary_factor") + else 1.0 + ) + head_dim = getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads + ) + dim = int(head_dim * partial_rotary_factor) + + short_factor = rope_scaling.get("short_factor") + if not isinstance(short_factor, list) and all( + isinstance(x, (int, float)) for x in short_factor + ): + logger.warning( + "`rope_scaling`'s short_factor field must be a list of numbers, got %s", + short_factor, + ) + if not len(short_factor) == dim // 2: + logger.warning( + "`rope_scaling`'s short_factor field must have length %s, got %s", + dim // 2, + len(short_factor), + ) + + long_factor = rope_scaling.get("long_factor") + if not isinstance(long_factor, list) and all( + isinstance(x, (int, float)) for x in long_factor + ): + logger.warning( + "`rope_scaling`'s long_factor field must be a list of numbers, got %s", + long_factor, + ) + if not len(long_factor) == dim // 2: + logger.warning( + "`rope_scaling`'s long_factor field must have length %s, got %s", + dim // 2, + len(long_factor), + ) + + # Handle Phi3 divergence: prefer the use of `attention_factor` and/or `factor` over + # `original_max_position_embeddings` to compute internal variables. The latter lives outside `rope_scaling` and is + # unique to longrope (= undesirable) + if hasattr(config, "original_max_position_embeddings"): + logger.warning_once( + "This model has set a `original_max_position_embeddings` field, to be used together with " + "`max_position_embeddings` to determine a scaling factor. Please set the `factor` field of `rope_scaling`" + "with this ratio instead -- we recommend the use of this field over `original_max_position_embeddings`, " + "as it is compatible with most model architectures." + ) + else: + factor = rope_scaling.get("factor") + if factor is None: + logger.warning("Missing required keys in `rope_scaling`: 'factor'") + elif not isinstance(factor, float) or factor < 1.0: + logger.warning( + "`rope_scaling`'s factor field must be a float >= 1, got %s", factor + ) + + attention_factor = rope_scaling.get("attention_factor") + if attention_factor is not None: + if not isinstance(attention_factor, float) or attention_factor < 0.0: + logger.warning( + "`rope_scaling`'s attention_factor field must be a float greater than 0, got %s", + attention_factor, + ) + + +def _validate_llama3_parameters( + config: PretrainedConfig, ignore_keys: set | None = None +): + rope_scaling = config.rope_scaling + rope_type = rope_scaling.get( + "rope_type", rope_scaling.get("type", None) + ) # BC: "rope_type" was originally "type" + required_keys = { + "rope_type", + "factor", + "original_max_position_embeddings", + "low_freq_factor", + "high_freq_factor", + } + received_keys = set(rope_scaling.keys()) + _check_received_keys( + rope_type, received_keys, required_keys, ignore_keys=ignore_keys + ) + + factor = rope_scaling["factor"] + if factor is None or not isinstance(factor, float) or factor < 1.0: + logger.warning( + "`rope_scaling`'s factor field must be a float >= 1, got %s", factor + ) + + low_freq_factor = rope_scaling["low_freq_factor"] + high_freq_factor = rope_scaling["high_freq_factor"] + if low_freq_factor is None or not isinstance(low_freq_factor, float): + logger.warning( + "`rope_scaling`'s low_freq_factor field must be a float, got %s", + low_freq_factor, + ) + if high_freq_factor is None or not isinstance(high_freq_factor, float): + logger.warning( + "`rope_scaling`'s high_freq_factor field must be a float, got %s", + high_freq_factor, + ) + if high_freq_factor <= low_freq_factor: + logger.warning( + "`rope_scaling`'s high_freq_factor field must be greater than low_freq_factor, got high_freq_factor=%s and low_freq_factor=%s", + high_freq_factor, + low_freq_factor, + ) + + original_max_position_embeddings = rope_scaling["original_max_position_embeddings"] + if original_max_position_embeddings is None or not isinstance( + original_max_position_embeddings, int + ): + logger.warning( + "`rope_scaling`'s original_max_position_embeddings field must be an integer, got %s", + original_max_position_embeddings, + ) + if original_max_position_embeddings >= config.max_position_embeddings: + logger.warning( + "`rope_scaling`'s original_max_position_embeddings field must be less than max_position_embeddings, got %s and max_position_embeddings=%s", + original_max_position_embeddings, + config.max_position_embeddings, + ) + + +# Like `ROPE_INIT_FUNCTIONS`, this validation function mapping can be dynamically updated for custom RoPE types. +ROPE_VALIDATION_FUNCTIONS = { + "default": _validate_default_rope_parameters, + "linear": _validate_linear_scaling_rope_parameters, + "dynamic": _validate_dynamic_scaling_rope_parameters, + "yarn": _validate_yarn_parameters, + "longrope": _validate_longrope_parameters, + "llama3": _validate_llama3_parameters, +} + + +def rope_config_validation(config: PretrainedConfig, ignore_keys: set | None = None): + """ + Validate the RoPE config arguments, given a `PretrainedConfig` object + """ + rope_scaling = getattr( + config, "rope_scaling", None + ) # not a default parameter in `PretrainedConfig` + if rope_scaling is None: + return + + # BC: "rope_type" was originally "type" + rope_type = rope_scaling.get("rope_type", rope_scaling.get("type", "default")) + validation_fn = ROPE_VALIDATION_FUNCTIONS.get(rope_type) + if validation_fn is not None: + validation_fn(config, ignore_keys=ignore_keys) + else: + logger.warning( + "Missing validation function mapping in `ROPE_VALIDATION_FUNCTIONS` for 'rope_type'='%s'", + rope_type, + ) diff --git a/python/tokenspeed/runtime/distributed/__init__.py b/python/tokenspeed/runtime/distributed/__init__.py new file mode 100755 index 0000000..2a5835d --- /dev/null +++ b/python/tokenspeed/runtime/distributed/__init__.py @@ -0,0 +1,26 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Distributed runtime public exports.""" + +from tokenspeed.runtime.distributed.comm_manager import CommManager +from tokenspeed.runtime.distributed.mapping import Mapping + +__all__ = ["CommManager", "Mapping"] diff --git a/python/tokenspeed/runtime/distributed/comm_backend/__init__.py b/python/tokenspeed/runtime/distributed/comm_backend/__init__.py new file mode 100644 index 0000000..a69cade --- /dev/null +++ b/python/tokenspeed/runtime/distributed/comm_backend/__init__.py @@ -0,0 +1,32 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from tokenspeed.runtime.distributed.comm_backend.base import CommBackend, Group +from tokenspeed.runtime.distributed.comm_backend.registry import ( + get_global_backend, + initialize_comm_backend, +) + +__all__ = [ + "CommBackend", + "Group", + "get_global_backend", + "initialize_comm_backend", +] diff --git a/python/tokenspeed/runtime/distributed/comm_backend/auto.py b/python/tokenspeed/runtime/distributed/comm_backend/auto.py new file mode 100644 index 0000000..82c21c7 --- /dev/null +++ b/python/tokenspeed/runtime/distributed/comm_backend/auto.py @@ -0,0 +1,134 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Auto backend: per-call strategy selection. + +Wraps NcclBackend and optionally CustomAllReduceBackend and +the fused all-reduce backend. For all_reduce, selects the lowest-latency +backend based on tensor size and hardware. For other ops, always uses +NCCL. +""" + +import torch + +from tokenspeed.runtime.distributed.comm_backend.base import CommBackend, Group +from tokenspeed.runtime.distributed.comm_backend.custom_allreduce import ( + CustomAllReduceBackend, +) +from tokenspeed.runtime.distributed.comm_backend.nccl import NcclBackend +from tokenspeed.runtime.distributed.comm_backend.triton_allreduce import ( + TritonAllReduceBackend, +) +from tokenspeed.runtime.distributed.comm_backend.triton_rsag import TritonRSAGBackend +from tokenspeed.runtime.distributed.comm_backend.trtllm_allreduce import ( + TrtllmAllReduceBackend, +) + + +class AutoBackend(CommBackend): + """Composite backend that selects the best strategy per call.""" + + def __init__(self): + self._nccl = NcclBackend() + self._custom_ar = CustomAllReduceBackend(fallback=self._nccl) + self._trtllm_ar = TrtllmAllReduceBackend(fallback=self._nccl) + self._triton_ar = TritonAllReduceBackend(fallback=self._nccl) + self._rsag = TritonRSAGBackend(fallback=self._nccl) + + @property + def nccl(self) -> NcclBackend: + return self._nccl + + @property + def custom_ar(self) -> CustomAllReduceBackend: + return self._custom_ar + + @property + def trtllm_ar(self) -> TrtllmAllReduceBackend: + return self._trtllm_ar + + def configure( + self, use_pynccl: bool = False, use_custom_allreduce: bool = False + ) -> None: + self._nccl.configure(use_pynccl=use_pynccl) + self._custom_ar.configure(use_custom_allreduce=use_custom_allreduce) + + # ---- Token-aware ops ---- + + def token_all_gather( + self, + tensor: torch.Tensor, + group: Group, + scattered_num_tokens: list[int], + ) -> torch.Tensor: + return self._rsag.token_all_gather(tensor, group, scattered_num_tokens) + + def token_reduce_scatter( + self, + tensor: torch.Tensor, + group: Group, + scattered_num_tokens: list[int], + ) -> torch.Tensor: + return self._rsag.token_reduce_scatter(tensor, group, scattered_num_tokens) + + # ---- Public CommBackend interface ---- + + def all_reduce(self, tensor: torch.Tensor, group: Group, op=None) -> torch.Tensor: + if self._custom_ar.has_custom_ar(group): + return self._custom_ar.all_reduce(tensor, group, op=op) + if self._trtllm_ar.has_trtllm_ar(group): + return self._trtllm_ar.all_reduce(tensor, group, op=op) + if self._triton_ar.can_run(tensor, group, op=op): + return self._triton_ar.all_reduce(tensor, group, op=op) + return self._nccl.all_reduce(tensor, group, op=op) + + def all_gather( + self, tensor: torch.Tensor, group: Group, dim: int = 0 + ) -> torch.Tensor: + if tensor.dim() == 2 and dim in (-1, tensor.dim() - 1): + return self._rsag.all_gather(tensor, group, dim) + + return self._nccl.all_gather(tensor, group, dim) + + def all_gather_into_tensor( + self, output: torch.Tensor, input: torch.Tensor, group: Group + ) -> None: + return self._nccl.all_gather_into_tensor(output, input, group) + + def reduce_scatter(self, tensor: torch.Tensor, group: Group) -> torch.Tensor: + return self._nccl.reduce_scatter(tensor, group) + + def all_to_all_single( + self, output: torch.Tensor, input: torch.Tensor, group: Group + ) -> None: + return self._nccl.all_to_all_single(output, input, group) + + def send(self, tensor: torch.Tensor, dst: int, group: Group) -> None: + return self._nccl.send(tensor, dst, group) + + def recv( + self, + size: torch.Size, + dtype: torch.dtype, + device: torch.device, + src: int, + group: Group, + ) -> torch.Tensor: + return self._nccl.recv(size, dtype, device, src, group) diff --git a/python/tokenspeed/runtime/distributed/comm_backend/base.py b/python/tokenspeed/runtime/distributed/comm_backend/base.py new file mode 100644 index 0000000..04807db --- /dev/null +++ b/python/tokenspeed/runtime/distributed/comm_backend/base.py @@ -0,0 +1,82 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Abstract base class for communication backends.""" + +from abc import ABC, abstractmethod + +import torch + +from tokenspeed.runtime.distributed.mapping import Group + + +class CommBackend(ABC): + """Interface that all communication backends must implement. + + All group parameters are tuples of global ranks, e.g. (0, 1, 2, 3). + Process groups are looked up from pg_manager, not created here. + """ + + # ---- Collective ops ---- + + @abstractmethod + def all_reduce( + self, tensor: torch.Tensor, group: Group, op=None + ) -> torch.Tensor: ... + + @abstractmethod + def all_gather( + self, tensor: torch.Tensor, group: Group, dim: int = 0 + ) -> torch.Tensor: ... + + @abstractmethod + def all_gather_into_tensor( + self, output: torch.Tensor, input: torch.Tensor, group: Group + ) -> None: ... + + @abstractmethod + def reduce_scatter(self, tensor: torch.Tensor, group: Group) -> torch.Tensor: ... + + @abstractmethod + def all_to_all_single( + self, output: torch.Tensor, input: torch.Tensor, group: Group + ) -> None: + """Even-split all_to_all. output and input must have same numel + divisible by len(group). + """ + ... + + # ---- Token-aware ops (uneven token distribution) ---- + + @abstractmethod + def token_all_gather( + self, + tensor: torch.Tensor, + group: Group, + scattered_num_tokens: list[int], + ) -> torch.Tensor: ... + + @abstractmethod + def token_reduce_scatter( + self, + tensor: torch.Tensor, + group: Group, + scattered_num_tokens: list[int], + ) -> torch.Tensor: ... diff --git a/python/tokenspeed/runtime/distributed/comm_backend/custom_allreduce.py b/python/tokenspeed/runtime/distributed/comm_backend/custom_allreduce.py new file mode 100644 index 0000000..aea2d63 --- /dev/null +++ b/python/tokenspeed/runtime/distributed/comm_backend/custom_allreduce.py @@ -0,0 +1,139 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Custom all-reduce backend using P2P GPU shared memory. + +Only supports all_reduce. Other ops delegate to a fallback backend. +""" + +from contextlib import nullcontext + +import torch + +from tokenspeed.runtime.distributed.comm_backend.base import CommBackend, Group + + +class CustomAllReduceBackend(CommBackend): + """Backend using custom P2P all-reduce (NVLink shared memory). + + Maintains per-group ca_comm in an internal registry, keyed by group tuple. + Falls back to the provided fallback backend for ops other than + all_reduce, or when the tensor is not eligible for custom AR. + """ + + def __init__(self, fallback: CommBackend): + self._fallback = fallback + self._resources = {} # group_tuple → {ca_comm} + self._use_custom_allreduce = False + + def configure(self, use_custom_allreduce: bool = False) -> None: + self._use_custom_allreduce = use_custom_allreduce + + def _get_or_create_resources(self, group: Group): + if group in self._resources: + return self._resources[group] + + ca_comm = None + if self._use_custom_allreduce and len(group) > 1: + try: + from tokenspeed.runtime.distributed.device_communicators.custom_all_reduce import ( + CustomAllreduce, + ) + from tokenspeed.runtime.distributed.process_group_manager import ( + process_group_manager as pg_manager, + ) + + gloo_group = pg_manager.get_process_group("gloo", group) + ca_comm = CustomAllreduce( + group=gloo_group, + device=torch.device(f"cuda:{torch.cuda.current_device()}"), + ) + except Exception: + ca_comm = None + + self._resources[group] = {"ca_comm": ca_comm} + return self._resources[group] + + def has_custom_ar(self, group: Group) -> bool: + if group not in self._resources: + return False + res = self._resources[group] + ca_comm = res["ca_comm"] + return ca_comm is not None and not ca_comm.disabled + + def capture(self, group: Group): + res = self._get_or_create_resources(group) + ca_comm = res["ca_comm"] + if ca_comm is None or ca_comm.disabled: + return nullcontext() + return ca_comm.capture() + + # ---- Public CommBackend interface ---- + + def all_reduce(self, tensor: torch.Tensor, group: Group, op=None) -> torch.Tensor: + if op is None: + op = torch.distributed.ReduceOp.SUM + res = self._get_or_create_resources(group) + ca_comm = res["ca_comm"] + if ( + op == torch.distributed.ReduceOp.SUM + and ca_comm is not None + and not ca_comm.disabled + and ca_comm.should_custom_ar(tensor) + ): + out = ca_comm.custom_all_reduce(tensor) + if out is None: + raise RuntimeError("custom all-reduce returned no output") + return out + return self._fallback.all_reduce(tensor, group, op=op) + + def all_gather( + self, tensor: torch.Tensor, group: Group, dim: int = 0 + ) -> torch.Tensor: + return self._fallback.all_gather(tensor, group, dim) + + def all_gather_into_tensor( + self, output: torch.Tensor, input: torch.Tensor, group: Group + ) -> None: + return self._fallback.all_gather_into_tensor(output, input, group) + + def reduce_scatter(self, tensor: torch.Tensor, group: Group) -> torch.Tensor: + return self._fallback.reduce_scatter(tensor, group) + + def all_to_all_single( + self, output: torch.Tensor, input: torch.Tensor, group: Group + ) -> None: + return self._fallback.all_to_all_single(output, input, group) + + def token_all_gather( + self, + tensor: torch.Tensor, + group: Group, + scattered_num_tokens: list[int], + ) -> torch.Tensor: + raise NotImplementedError("Use AutoBackend for token-aware ops") + + def token_reduce_scatter( + self, + tensor: torch.Tensor, + group: Group, + scattered_num_tokens: list[int], + ) -> torch.Tensor: + raise NotImplementedError("Use AutoBackend for token-aware ops") diff --git a/python/tokenspeed/runtime/distributed/comm_backend/nccl.py b/python/tokenspeed/runtime/distributed/comm_backend/nccl.py new file mode 100644 index 0000000..a32d129 --- /dev/null +++ b/python/tokenspeed/runtime/distributed/comm_backend/nccl.py @@ -0,0 +1,258 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""NCCL communication backend. + +Looks up pre-created process groups from pg_manager. Optionally uses +PyNccl communicators for better performance. Supports torch.compile +via custom ops. +""" + +import torch +import torch.distributed + +from tokenspeed.runtime.distributed.comm_backend.base import CommBackend, Group + + +class NcclBackend(CommBackend): + """Backend using NCCL via PyNccl or torch.distributed. + + Caches per-group resources (process group handle, PyNccl comm) + keyed by group tuple. Process groups are looked up from pg_manager + on first use. + """ + + def __init__(self): + self._resources = {} # group_tuple → {pynccl_comm, device_group, world_size} + self._use_pynccl = False + + def configure(self, use_pynccl: bool = False) -> None: + self._use_pynccl = use_pynccl + + def _get_or_create_resources(self, group: Group): + if group in self._resources: + return self._resources[group] + + from tokenspeed.runtime.distributed.process_group_manager import ( + process_group_manager as pg_manager, + ) + + device_group = pg_manager.get_process_group("nccl", group) + world_size = len(group) + + pynccl_comm = None + if self._use_pynccl and world_size > 1: + try: + from tokenspeed.runtime.distributed.device_communicators.pynccl import ( + PyNcclCommunicator, + ) + + gloo_group = pg_manager.get_process_group("gloo", group) + pynccl_comm = PyNcclCommunicator( + group=gloo_group, + device=torch.device(f"cuda:{torch.cuda.current_device()}"), + ) + except Exception: + pynccl_comm = None + + self._resources[group] = { + "pynccl_comm": pynccl_comm, + "device_group": device_group, + "world_size": world_size, + } + return self._resources[group] + + # ---- Public CommBackend interface ---- + + def all_reduce(self, tensor: torch.Tensor, group: Group, op=None) -> torch.Tensor: + res = self._get_or_create_resources(group) + if res["world_size"] == 1: + return tensor + if op is None: + op = torch.distributed.ReduceOp.SUM + pynccl = res["pynccl_comm"] + if pynccl is not None and not pynccl.disabled: + pynccl.all_reduce(tensor, op=op) + else: + torch.distributed.all_reduce(tensor, op=op, group=res["device_group"]) + return tensor + + def all_gather( + self, tensor: torch.Tensor, group: Group, dim: int = 0 + ) -> torch.Tensor: + res = self._get_or_create_resources(group) + ws = res["world_size"] + if ws == 1: + return tensor + if dim < 0: + dim += tensor.dim() + + input_size = tensor.size() + output_size = (input_size[0] * ws,) + input_size[1:] + output_tensor = torch.empty( + output_size, dtype=tensor.dtype, device=tensor.device + ) + + self.all_gather_into_tensor(output_tensor, tensor, group) + + output_tensor = output_tensor.reshape((ws,) + input_size) + output_tensor = output_tensor.movedim(0, dim) + output_tensor = output_tensor.reshape( + input_size[:dim] + (ws * input_size[dim],) + input_size[dim + 1 :] + ) + return output_tensor + + def all_gather_into_tensor( + self, output: torch.Tensor, input: torch.Tensor, group: Group + ) -> None: + res = self._get_or_create_resources(group) + pynccl = res["pynccl_comm"] + if pynccl is not None and not pynccl.disabled: + pynccl.all_gather(output, input) + else: + torch.distributed.all_gather_into_tensor( + output, input, group=res["device_group"] + ) + + def all_to_all_single( + self, output: torch.Tensor, input: torch.Tensor, group: Group + ) -> None: + res = self._get_or_create_resources(group) + ws = res["world_size"] + if ws == 1: + output.copy_(input) + return + # PyNccl has no all_to_all wrapper + torch.distributed.all_to_all_single(output, input, group=res["device_group"]) + + def reduce_scatter(self, tensor: torch.Tensor, group: Group) -> torch.Tensor: + res = self._get_or_create_resources(group) + ws = res["world_size"] + if ws == 1: + return tensor + input_size = tuple(tensor.size()) + output_tensor = torch.empty( + (input_size[0] // ws,) + input_size[1:], + dtype=tensor.dtype, + device=tensor.device, + ) + pynccl = res["pynccl_comm"] + if pynccl is not None and not pynccl.disabled: + pynccl.reduce_scatter(output_tensor, tensor) + else: + torch.distributed.reduce_scatter_tensor( + output_tensor, tensor, group=res["device_group"] + ) + return output_tensor + + def send(self, tensor: torch.Tensor, dst: int, group: Group) -> None: + res = self._get_or_create_resources(group) + pynccl = res["pynccl_comm"] + if pynccl is not None and not pynccl.disabled: + pynccl.send(tensor, dst) + else: + torch.distributed.send(tensor, group[dst], group=res["device_group"]) + + def recv( + self, + size: torch.Size, + dtype: torch.dtype, + device: torch.device, + src: int, + group: Group, + ) -> torch.Tensor: + res = self._get_or_create_resources(group) + tensor = torch.empty(size, dtype=dtype, device=device) + pynccl = res["pynccl_comm"] + if pynccl is not None and not pynccl.disabled: + pynccl.recv(tensor, src) + else: + torch.distributed.recv(tensor, group[src], group=res["device_group"]) + return tensor + + def token_all_gather( + self, + tensor: torch.Tensor, + group: Group, + scattered_num_tokens: list[int], + ) -> torch.Tensor: + """NCCL token_all_gather with padding for uneven token distribution. + + Pads each rank's slice to max_tokens rows, all-gathers, then strips padding. + """ + tp_size = len(scattered_num_tokens) + max_tokens = max(scattered_num_tokens) + hidden = tensor.size(-1) + + local_tokens = tensor.size(0) + if local_tokens < max_tokens: + pad = torch.zeros( + max_tokens - local_tokens, + hidden, + dtype=tensor.dtype, + device=tensor.device, + ) + padded = torch.cat([tensor, pad], dim=0) + else: + padded = tensor + + output = torch.empty( + tp_size * max_tokens, hidden, dtype=tensor.dtype, device=tensor.device + ) + self.all_gather_into_tensor(output, padded.contiguous(), group) + + chunks = [] + for i, n in enumerate(scattered_num_tokens): + chunks.append(output[i * max_tokens : i * max_tokens + n]) + return torch.cat(chunks, dim=0) + + def token_reduce_scatter( + self, + tensor: torch.Tensor, + group: Group, + scattered_num_tokens: list[int], + ) -> torch.Tensor: + """NCCL token_reduce_scatter with padding for uneven token distribution. + + Pads the gathered tensor to a uniform layout, reduce-scatters, then strips padding. + """ + tp_size = len(scattered_num_tokens) + max_tokens = max(scattered_num_tokens) + hidden = tensor.size(-1) + + padded_input = torch.zeros( + tp_size * max_tokens, hidden, dtype=tensor.dtype, device=tensor.device + ) + offset = 0 + for i, n in enumerate(scattered_num_tokens): + padded_input[i * max_tokens : i * max_tokens + n].copy_( + tensor[offset : offset + n] + ) + offset += n + + output = torch.empty( + max_tokens, hidden, dtype=tensor.dtype, device=tensor.device + ) + res = self._get_or_create_resources(group) + torch.distributed.reduce_scatter_tensor( + output, padded_input.contiguous(), group=res["device_group"] + ) + rank = group.index(torch.distributed.get_rank()) + return output[: scattered_num_tokens[rank]].contiguous() diff --git a/python/tokenspeed/runtime/distributed/comm_backend/registry.py b/python/tokenspeed/runtime/distributed/comm_backend/registry.py new file mode 100644 index 0000000..3c15f68 --- /dev/null +++ b/python/tokenspeed/runtime/distributed/comm_backend/registry.py @@ -0,0 +1,53 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Global communication backend management. + +Provides global singleton instances for CommBackend and TritonRSAGBackend. +""" + +from tokenspeed.runtime.distributed.comm_backend.base import CommBackend + +_global_backend: CommBackend | None = None + + +def initialize_comm_backend( + use_pynccl: bool = False, + use_custom_allreduce: bool = False, +) -> CommBackend: + """Create and configure the global communication backend.""" + global _global_backend + + from tokenspeed.runtime.distributed.comm_backend.auto import AutoBackend + + _global_backend = AutoBackend() + _global_backend.configure( + use_pynccl=use_pynccl, + use_custom_allreduce=use_custom_allreduce, + ) + return _global_backend + + +def get_global_backend() -> CommBackend: + """Get the global CommBackend, creating AutoBackend if not initialized.""" + global _global_backend + if _global_backend is None: + initialize_comm_backend() + return _global_backend diff --git a/python/tokenspeed/runtime/distributed/comm_backend/triton_allreduce.py b/python/tokenspeed/runtime/distributed/comm_backend/triton_allreduce.py new file mode 100644 index 0000000..d85e5c4 --- /dev/null +++ b/python/tokenspeed/runtime/distributed/comm_backend/triton_allreduce.py @@ -0,0 +1,115 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Triton all-reduce backend for latency-sensitive small AMD tensors.""" + +import torch +import torch.distributed as dist +from tokenspeed_kernel.ops.communication.triton import ( + all_reduce, + all_reduce_can_run, + create_state, +) + +from tokenspeed.runtime.distributed.comm_backend.base import CommBackend, Group +from tokenspeed.runtime.distributed.process_group_manager import ( + process_group_manager as pg_manager, +) + + +class TritonAllReduceBackend(CommBackend): + def __init__(self, fallback: CommBackend): + self._fallback = fallback + self._instances = {} + self._max_bytes = 512 * 1024 + self._max_numel = ( + self._max_bytes // torch.empty((), dtype=torch.bfloat16).element_size() + ) + + def _get_or_create(self, group: Group): + if group in self._instances: + return self._instances[group] + + state = create_state( + group=pg_manager.get_process_group("nccl", group), + rank_in_group=group.index(dist.get_rank()), + max_numel=self._max_numel, + device=torch.device(f"cuda:{torch.cuda.current_device()}"), + ) + self._instances[group] = state + return state + + def can_run(self, tensor: torch.Tensor, group: Group, op=None) -> bool: + if len(group) <= 1: + return False + if op is None: + op = torch.distributed.ReduceOp.SUM + if not ( + op == torch.distributed.ReduceOp.SUM + and tensor.is_cuda + and tensor.is_contiguous() + and tensor.dtype == torch.bfloat16 + and 0 < tensor.numel() <= self._max_numel + ): + return False + try: + return all_reduce_can_run(self._get_or_create(group), tensor, op=op) + except Exception: + return False + + def all_reduce(self, tensor: torch.Tensor, group: Group, op=None) -> torch.Tensor: + state = self._get_or_create(group) + if all_reduce_can_run(state, tensor, op=op): + return all_reduce(state, tensor, op=op) + return self._fallback.all_reduce(tensor, group, op=op) + + def all_gather( + self, tensor: torch.Tensor, group: Group, dim: int = 0 + ) -> torch.Tensor: + return self._fallback.all_gather(tensor, group, dim) + + def all_gather_into_tensor( + self, output: torch.Tensor, input: torch.Tensor, group: Group + ) -> None: + return self._fallback.all_gather_into_tensor(output, input, group) + + def reduce_scatter(self, tensor: torch.Tensor, group: Group) -> torch.Tensor: + return self._fallback.reduce_scatter(tensor, group) + + def all_to_all_single( + self, output: torch.Tensor, input: torch.Tensor, group: Group + ) -> None: + return self._fallback.all_to_all_single(output, input, group) + + def token_all_gather( + self, + tensor: torch.Tensor, + group: Group, + scattered_num_tokens: list[int], + ) -> torch.Tensor: + raise NotImplementedError("Use AutoBackend for token-aware ops") + + def token_reduce_scatter( + self, + tensor: torch.Tensor, + group: Group, + scattered_num_tokens: list[int], + ) -> torch.Tensor: + raise NotImplementedError("Use AutoBackend for token-aware ops") diff --git a/python/tokenspeed/runtime/distributed/comm_backend/triton_rsag.py b/python/tokenspeed/runtime/distributed/comm_backend/triton_rsag.py new file mode 100644 index 0000000..57b777d --- /dev/null +++ b/python/tokenspeed/runtime/distributed/comm_backend/triton_rsag.py @@ -0,0 +1,143 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""TritonRSAG communication backend for token-aware all_gather / reduce_scatter. + +Handles uneven token distribution across ranks using Triton RS/AG state. +Lazily creates and caches Triton RS/AG state keyed by (group_tuple, hidden_size). +""" + +import torch +import torch.distributed as dist +from tokenspeed_kernel.ops.communication.triton import ( + all_gather, + all_gather_inner, + create_state, + reduce_scatter, +) +from tokenspeed_kernel.platform import current_platform + +from tokenspeed.runtime.distributed.comm_backend.base import CommBackend, Group +from tokenspeed.runtime.distributed.process_group_manager import ( + process_group_manager as pg_manager, +) +from tokenspeed.runtime.utils import ceil_div +from tokenspeed.runtime.utils.env import global_server_args_dict + + +class TritonRSAGBackend: + """Backend using TritonRSAG for token-aware reduce_scatter / all_gather. + + Unlike NCCL backends, TritonRSAG handles uneven token distribution + across ranks (scattered tokens). Each instance is specific to a + (group, hidden_size) pair because RSAG pre-allocates buffers. + """ + + def __init__(self, fallback: CommBackend): + self._fallback = fallback + # (group_tuple, hidden_size) -> Triton RS/AG state + self._instances = {} + + def _get_or_create(self, group: Group, hidden_size: int): + key = (group, hidden_size) + if key in self._instances: + return self._instances[key] + + max_num_tokens = self._get_max_num_gathered_tokens() + state = create_state( + group=pg_manager.get_process_group("nccl", group), + rank_in_group=group.index(dist.get_rank()), + max_tokens=max_num_tokens, + hidden_size=hidden_size, + ) + self._instances[key] = state + return state + + def all_gather( + self, + tensor: torch.Tensor, + group: Group, + dim: int = 0, + ) -> torch.Tensor: + if tensor.dim() != 2: + return self._fallback.all_gather(tensor, group=group, dim=dim) + + if dim == 0: + return self.token_all_gather( + tensor, + group=group, + scattered_num_tokens=[tensor.size(0)] * len(group), + ) + + if ( + current_platform().is_nvidia + and dim in (-1, tensor.dim() - 1) + and tensor.dtype == torch.bfloat16 + ): + hidden_size = tensor.size(-1) * len(group) + state = self._get_or_create(group, hidden_size) + return all_gather_inner( + state, + tensor, + tp_hidden_dim=hidden_size, + skip_entry_sync=False, + safe=False, + ) + + return self._fallback.all_gather(tensor, group=group, dim=dim) + + def token_all_gather( + self, + tensor: torch.Tensor, + group: Group, + scattered_num_tokens: list[int], + ) -> torch.Tensor: + state = self._get_or_create(group, tensor.size(-1)) + return all_gather(state, tensor, token_list_in_group=scattered_num_tokens) + + def token_reduce_scatter( + self, + tensor: torch.Tensor, + group: Group, + scattered_num_tokens: list[int], + ) -> torch.Tensor: + state = self._get_or_create(group, tensor.size(-1)) + return reduce_scatter(state, tensor, token_list_in_group=scattered_num_tokens) + + def _get_max_num_gathered_tokens(self): + """Compute max buffer size for TritonRSAG. + + global_server_args_dict read is intentional — this is one-time RSAG buffer + init infrastructure. Passing mapping through all signatures would be too invasive. + """ + mapping = global_server_args_dict["mapping"] + chunked_prefill_size = global_server_args_dict["chunked_prefill_size"] + max_prefill_tokens = global_server_args_dict["max_prefill_tokens"] + max_model_len = global_server_args_dict["max_model_len"] + if chunked_prefill_size > 0: + max_attn_tp_num_tokens = chunked_prefill_size + else: + max_attn_tp_num_tokens = max_prefill_tokens + max_model_len + max_scattered_num_tokens = ceil_div( + max_attn_tp_num_tokens, mapping.attn.tp_size + ) + return max_scattered_num_tokens * max( + mapping.attn.tp_size, mapping.dense.tp_size, mapping.moe.tp_ep_size + ) diff --git a/python/tokenspeed/runtime/distributed/comm_backend/trtllm_allreduce.py b/python/tokenspeed/runtime/distributed/comm_backend/trtllm_allreduce.py new file mode 100644 index 0000000..ec25ffd --- /dev/null +++ b/python/tokenspeed/runtime/distributed/comm_backend/trtllm_allreduce.py @@ -0,0 +1,213 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Lamport 1-shot all-reduce backend. + +Uses an IPC workspace with Lamport barriers and shared memory for low-latency +all-reduce on small tensors. Falls back to a provided fallback backend for +large tensors or unsupported ops. + +The workspace is created once per group via ``configure_group`` and +reused for every subsequent ``all_reduce`` on that group. +""" + +import torch +from tokenspeed_kernel.ops.communication.trtllm import ( + AllReduceFusionPattern, + trtllm_allreduce_fusion, + trtllm_create_ipc_workspace_for_all_reduce_fusion, +) +from tokenspeed_kernel.platform import current_platform + +from tokenspeed.runtime.distributed.comm_backend.base import CommBackend, Group + +_MAX_ONESHOT_BYTES = 2 * 1024 * 1024 + + +class TrtllmAllReduceBackend(CommBackend): + """Backend using Lamport 1-shot all-reduce. + + Keyed per-group: each group gets its own IPC workspace so handles + are never reused across groups. Only ``all_reduce`` (SUM) is + accelerated; every other op delegates to *fallback*. + """ + + def __init__(self, fallback: CommBackend): + self._fallback = fallback + self._resources = {} # group_tuple → {workspace, rank, world_size} + + def _load_comm(self): + return current_platform().is_nvidia + + # ------------------------------------------------------------------ + # Group configuration + # ------------------------------------------------------------------ + + def configure_group( + self, + rank: int, + group: Group, + max_token_num: int, + hidden_dim: int, + use_fp32_lamport: bool = False, + ) -> bool: + """Create IPC workspace for *group*. Returns True on success.""" + if group in self._resources: + return True + + if not self._load_comm(): + return False + + try: + + from tokenspeed.runtime.distributed.process_group_manager import ( + process_group_manager as pg_manager, + ) + + device_group = pg_manager.get_process_group("nccl", group) + + ipc_handles, workspace_tensor = ( + trtllm_create_ipc_workspace_for_all_reduce_fusion( + rank, + len(group), + max_token_num, + hidden_dim, + group=device_group, + use_fp32_lamport=use_fp32_lamport, + ) + ) + + self._resources[group] = { + "ipc_handles": ipc_handles, + "workspace": workspace_tensor, + "rank": rank, + "world_size": len(group), + "max_token_num": max_token_num, + "hidden_dim": hidden_dim, + "device_group": device_group, + } + + return True + + except Exception: + + return False + + def has_trtllm_ar(self, group: Group) -> bool: + return group in self._resources + + # ------------------------------------------------------------------ + # CommBackend interface + # ------------------------------------------------------------------ + + def all_reduce(self, tensor: torch.Tensor, group: Group, op=None) -> torch.Tensor: + + if op is None: + op = torch.distributed.ReduceOp.SUM + + res = self._resources.get(group) + + if ( + res is not None + and op == torch.distributed.ReduceOp.SUM + and tensor.numel() * tensor.element_size() <= _MAX_ONESHOT_BYTES + ): + + result = self._lamport_allreduce(tensor, res) + + if result is not None: + return result + + return self._fallback.all_reduce(tensor, group, op=op) + + def _lamport_allreduce( + self, tensor: torch.Tensor, res: dict + ) -> torch.Tensor | None: + """Run the Lamport 1-shot kernel, return None on failure.""" + orig_shape = tensor.shape + + # The fused kernel expects 2D [token_num, hidden_dim]. + if tensor.dim() == 1: + tensor_2d = tensor.unsqueeze(0) + elif tensor.dim() > 2: + tensor_2d = tensor.reshape(-1, tensor.shape[-1]) + else: + tensor_2d = tensor + + token_num, hidden_dim = tensor_2d.shape + if hidden_dim > res["hidden_dim"] or token_num > res["max_token_num"]: + return None + + from tokenspeed.runtime.utils.pdl import pdl_enabled + + allreduce_out = torch.empty_like(tensor_2d) + + trtllm_allreduce_fusion( + allreduce_in=tensor_2d, + world_size=res["world_size"], + world_rank=res["rank"], + token_num=token_num, + hidden_dim=hidden_dim, + workspace_ptrs=res["workspace"], + launch_with_pdl=pdl_enabled(), + use_oneshot=True, + trigger_completion_at_end=True, + fp32_acc=False, + pattern_code=AllReduceFusionPattern.kAllReduce, + allreduce_out=allreduce_out, + ) + + return allreduce_out.view(orig_shape) + + # ---- Delegate everything else to fallback ---- + + def all_gather( + self, tensor: torch.Tensor, group: Group, dim: int = 0 + ) -> torch.Tensor: + return self._fallback.all_gather(tensor, group, dim) + + def all_gather_into_tensor( + self, output: torch.Tensor, input: torch.Tensor, group: Group + ) -> None: + return self._fallback.all_gather_into_tensor(output, input, group) + + def reduce_scatter(self, tensor: torch.Tensor, group: Group) -> torch.Tensor: + return self._fallback.reduce_scatter(tensor, group) + + def all_to_all_single( + self, output: torch.Tensor, input: torch.Tensor, group: Group + ) -> None: + return self._fallback.all_to_all_single(output, input, group) + + def token_all_gather( + self, + tensor: torch.Tensor, + group: Group, + scattered_num_tokens: list[int], + ) -> torch.Tensor: + raise NotImplementedError("Use AutoBackend for token-aware ops") + + def token_reduce_scatter( + self, + tensor: torch.Tensor, + group: Group, + scattered_num_tokens: list[int], + ) -> torch.Tensor: + raise NotImplementedError("Use AutoBackend for token-aware ops") diff --git a/python/tokenspeed/runtime/distributed/comm_manager.py b/python/tokenspeed/runtime/distributed/comm_manager.py new file mode 100755 index 0000000..9641d17 --- /dev/null +++ b/python/tokenspeed/runtime/distributed/comm_manager.py @@ -0,0 +1,371 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import torch + +from tokenspeed.runtime.distributed.comm_ops import ( + all_reduce, + token_all_gather, + token_reduce_scatter, +) +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.execution.context import ForwardContext + + +class CommManager: + """Manages communication patterns (all_reduce vs RSAG) for each decoder layer.""" + + def __init__( + self, + mapping: Mapping, + layer_id: int, + is_moe: bool, + prev_is_moe: bool, + input_layernorm: torch.nn.Module | None = None, + post_attn_layernorm: torch.nn.Module | None = None, + ) -> None: + self.mapping = mapping + self.layer_id = layer_id + self.is_moe = is_moe + self.prev_is_moe = prev_is_moe + self.input_layernorm = input_layernorm + self.post_attn_layernorm = post_attn_layernorm + + # ---- Scattered token counts ---- + + @staticmethod + def _scatter_count(num_tokens: int, tp_size: int) -> list[int]: + base, remainder = divmod(num_tokens, tp_size) + return [base + 1] * remainder + [base] * (tp_size - remainder) + + def get_num_tokens(self, ctx: ForwardContext): + scattered = self.scattered_num_tokens(ctx) + return sum(scattered), max(scattered) + + def scattered_num_tokens(self, ctx: ForwardContext) -> list[int]: + global_counts = ( + ctx.collective_global_num_tokens + if ctx.collective_global_num_tokens is not None + else ctx.global_num_tokens + ) + if global_counts is not None: + scattered = [] + for attn_dp_rank in range(self.mapping.attn.dp_size): + # global_counts is indexed by global rank with dp stride + # tp_size * cp_size; cp peers report the same count. + num_tokens = global_counts[ + attn_dp_rank * self.mapping.attn.tp_size * self.mapping.attn.cp_size + ] + scattered.extend( + self._scatter_count(num_tokens, self.mapping.attn.tp_size) + ) + return scattered + num_tokens = ( + ctx.collective_num_tokens + if ctx.collective_num_tokens is not None + else ctx.input_num_tokens + ) + return self._scatter_count(num_tokens, self.mapping.attn.tp_size) + + def attn_tp_group_scattered_num_tokens(self, ctx: ForwardContext) -> list[int]: + start = self.mapping.attn.tp_size * self.mapping.attn.dp_rank + end = start + self.mapping.attn.tp_size + return self.scattered_num_tokens(ctx)[start:end] + + def dense_tp_group_scattered_num_tokens(self, ctx: ForwardContext) -> list[int]: + start = self.mapping.dense.tp_size * self.mapping.dense.dp_rank + end = start + self.mapping.dense.tp_size + return self.scattered_num_tokens(ctx)[start:end] + + def moe_tp_ep_group_scattered_num_tokens(self, ctx: ForwardContext) -> list[int]: + tp_ep_size = self.mapping.moe.tp_ep_size + global_counts = ( + ctx.collective_global_num_tokens + if ctx.collective_global_num_tokens is not None + else ctx.global_num_tokens + ) + # Without DP, all ranks share the batch and the scattered table needs + # no global metadata, so the lookup below stays valid. + if global_counts is not None or not self.mapping.attn.has_dp: + # After post_attn_comm reduce-scatter, each rank holds its + # scattered share of its attn dp group's tokens, not the raw + # global count; MoE collectives must size from those rows. + scattered = self.scattered_num_tokens(ctx) + return [ + scattered[self.mapping.attn.scatter_index(rank)] + for rank in self.mapping.moe.tp_ep_group + ] + # With DP but no gathered metadata, other dp groups' counts are + # unknown; only the local rank's contribution can be reported. + num_tokens = ( + ctx.collective_num_tokens + if ctx.collective_num_tokens is not None + else ctx.input_num_tokens + ) + result = [0] * tp_ep_size + result[self.mapping.moe.tp_ep_rank] = num_tokens + return result + + # ---- Communication patterns ---- + + def use_all_reduce(self, is_moe: bool): + if is_moe: + return self.mapping.attn.tp_size == self.mapping.moe.tp_ep_size + return self.mapping.attn.tp_size == self.mapping.dense.tp_size + + def pre_attn_comm(self, hidden_states: torch.Tensor, ctx: ForwardContext): + if self.layer_id == 0: + return hidden_states + + if not self.mapping.has_attn_tp: + return hidden_states + + if self.use_all_reduce(self.prev_is_moe): + return hidden_states + + return token_all_gather( + hidden_states, + group=self.mapping.attn.tp_group, + scattered_num_tokens=self.attn_tp_group_scattered_num_tokens(ctx), + ) + + def gather_residual(self, residual: torch.Tensor, ctx: ForwardContext): + """All-gather a residual left scattered by the previous layer's RSAG + path (e.g. for aux hidden capture); no-op when rows are already full. + + Mirrors the pre_attn_comm gather conditions. + """ + if self.layer_id == 0: + return residual + if not self.mapping.has_attn_tp: + return residual + if self.use_all_reduce(self.prev_is_moe): + return residual + return token_all_gather( + residual, + group=self.mapping.attn.tp_group, + scattered_num_tokens=self.attn_tp_group_scattered_num_tokens(ctx), + ) + + def post_attn_comm( + self, hidden_states: torch.Tensor, residual: torch.Tensor, ctx: ForwardContext + ): + if not self.mapping.has_attn_tp: + return hidden_states, residual + + if self.use_all_reduce(self.is_moe): + hidden_states = all_reduce(hidden_states, self.mapping.attn.tp_group) + # The output residual is expected to have attn_tp_num_tokens. + # For first layer, the input residual has attn_tp_num_tokens. + # Otherwise, if this layer experiences a RSAG -> AR switch, residual needs allgather. + if self.layer_id > 0 and not self.use_all_reduce(self.prev_is_moe): + residual = token_all_gather( + residual, + group=self.mapping.attn.tp_group, + scattered_num_tokens=self.attn_tp_group_scattered_num_tokens(ctx), + ) + else: + token_list = self.attn_tp_group_scattered_num_tokens(ctx) + hidden_states = token_reduce_scatter( + hidden_states, + group=self.mapping.attn.tp_group, + scattered_num_tokens=token_list, + ) + # The output residual is expected to have scattered_num_tokens. + # For first layer, the input residual has attn_tp_num_tokens, so needs slice. + # Otherwise, if this layer experiences a AR -> RSAG switch, residual needs slice. + if self.layer_id == 0 or self.use_all_reduce(self.prev_is_moe): + offset = sum(token_list[: self.mapping.attn.tp_rank]) + residual = residual[offset : offset + hidden_states.size(0)] + + return hidden_states, residual + + def pre_mlp_comm(self, hidden_states: torch.Tensor, ctx: ForwardContext): + if self.is_moe: + return self.pre_moe_comm(hidden_states, ctx) + else: + return self.pre_dense_comm(hidden_states, ctx) + + def pre_dense_comm(self, hidden_states: torch.Tensor, ctx: ForwardContext): + if not self.mapping.dense.has_tp: + return hidden_states + + if self.use_all_reduce(is_moe=False): + return hidden_states + + return token_all_gather( + hidden_states, + group=self.mapping.dense.tp_group, + scattered_num_tokens=self.dense_tp_group_scattered_num_tokens(ctx), + ) + + def pre_moe_comm(self, hidden_states: torch.Tensor, ctx: ForwardContext): + if not self.mapping.moe.has_tp_ep: + return hidden_states + + if self.use_all_reduce(is_moe=True): + return hidden_states + + return token_all_gather( + hidden_states, + group=self.mapping.moe.tp_ep_group, + scattered_num_tokens=self.moe_tp_ep_group_scattered_num_tokens(ctx), + ) + + def post_mlp_comm( + self, hidden_states: torch.Tensor, residual: torch.Tensor, ctx: ForwardContext + ): + if self.is_moe: + return self.post_moe_comm(hidden_states, residual, ctx) + else: + return self.post_dense_comm(hidden_states, residual, ctx) + + def post_dense_comm( + self, hidden_states: torch.Tensor, residual: torch.Tensor, ctx: ForwardContext + ): + if not self.mapping.dense.has_tp: + return hidden_states, residual + + if self.use_all_reduce(is_moe=False): + hidden_states = all_reduce(hidden_states, self.mapping.dense.tp_group) + return hidden_states, residual + hidden_states = token_reduce_scatter( + hidden_states, + group=self.mapping.dense.tp_group, + scattered_num_tokens=self.dense_tp_group_scattered_num_tokens(ctx), + ) + return hidden_states, residual + + def post_moe_comm( + self, hidden_states: torch.Tensor, residual: torch.Tensor, ctx: ForwardContext + ): + if not self.mapping.moe.has_tp_ep: + return hidden_states, residual + + if self.use_all_reduce(is_moe=True): + hidden_states = all_reduce(hidden_states, self.mapping.moe.tp_ep_group) + return hidden_states, residual + hidden_states = token_reduce_scatter( + hidden_states, + group=self.mapping.moe.tp_ep_group, + scattered_num_tokens=self.moe_tp_ep_group_scattered_num_tokens(ctx), + ) + return hidden_states, residual + + def post_final_norm_comm( + self, hidden_states: torch.Tensor, residual: torch.Tensor, ctx: ForwardContext + ): + if not self.mapping.has_attn_tp: + return hidden_states, residual + if self.use_all_reduce(self.is_moe): + return hidden_states, residual + hidden_states = token_all_gather( + hidden_states, + group=self.mapping.attn.tp_group, + scattered_num_tokens=self.attn_tp_group_scattered_num_tokens(ctx), + ) + return hidden_states, residual + + # ---- Fused allreduce+norm ---- + + def use_all_reduce_norm_fusion(self) -> bool: + from tokenspeed.runtime.utils.env import global_server_args_dict + + return ( + self.use_all_reduce(self.is_moe) + and self.mapping.has_attn_tp + and global_server_args_dict.get("enable_allreduce_fusion", False) + ) + + def should_fuse(self, num_tokens: int) -> bool: + from tokenspeed.runtime.utils.env import global_server_args_dict + + return ( + self.use_all_reduce_norm_fusion() + and num_tokens > 0 + and num_tokens <= global_server_args_dict["comm_fusion_max_num_tokens"] + ) + + def input_reduce_norm( + self, hidden_states: torch.Tensor, residual: torch.Tensor | None + ): + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + elif self.should_fuse(hidden_states.shape[0]): + hidden_states, residual, *_ = ( + self.input_layernorm.forward_with_allreduce_fusion( + self.mapping.attn.tp_rank, + self.mapping.attn.tp_group, + hidden_states, + residual, + ) + ) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + return hidden_states, residual + + def post_attn_reduce_norm( + self, hidden_states: torch.Tensor, residual: torch.Tensor, ctx: ForwardContext + ): + if self.should_fuse(hidden_states.shape[0]): + hidden_states, residual, *_ = ( + self.post_attn_layernorm.forward_with_allreduce_fusion( + self.mapping.attn.tp_rank, + self.mapping.attn.tp_group, + hidden_states, + residual, + ) + ) + else: + hidden_states, residual = self.post_attn_comm(hidden_states, residual, ctx) + hidden_states, residual = self.post_attn_layernorm(hidden_states, residual) + return hidden_states, residual + + def post_mlp_fused( + self, hidden_states: torch.Tensor, residual: torch.Tensor, ctx: ForwardContext + ): + if not self.should_fuse(hidden_states.shape[0]): + hidden_states, residual = self.post_mlp_comm(hidden_states, residual, ctx) + return hidden_states, residual + + def final_norm( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor, + ctx: ForwardContext, + norm: torch.nn.Module, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + + if ctx.forward_mode.is_idle(): + return hidden_states, None + + if self.should_fuse(hidden_states.shape[0]): + hidden_states, residual_out, *_ = norm.forward_with_allreduce_fusion( + self.mapping.attn.tp_rank, + self.mapping.attn.tp_group, + hidden_states, + residual, + ) + else: + hidden_states, residual_out = norm(hidden_states, residual) + hidden_states, _ = self.post_final_norm_comm(hidden_states, residual, ctx) + + return hidden_states, residual_out diff --git a/python/tokenspeed/runtime/distributed/comm_ops.py b/python/tokenspeed/runtime/distributed/comm_ops.py new file mode 100644 index 0000000..7cf2e65 --- /dev/null +++ b/python/tokenspeed/runtime/distributed/comm_ops.py @@ -0,0 +1,319 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Communication ops for distributed communication. + +All ops require explicit group (tuple of ranks) and rank parameters. +Groups are looked up from pg_manager internally via comm_backend. +""" + +from dataclasses import dataclass +from enum import IntEnum + +import torch +import torch.distributed +from tokenspeed_kernel.ops.communication.trtllm import ( + allgather_dual_rmsnorm, + allreduce_residual_rmsnorm, + reducescatter_residual_rmsnorm, +) + +from tokenspeed.runtime.distributed.comm_backend import ( + CommBackend, + Group, + get_global_backend, +) +from tokenspeed.runtime.distributed.process_group_manager import ( + process_group_manager as pg_manager, +) +from tokenspeed.runtime.utils.pdl import pdl_enabled + + +def _get_process_group(group: Group): + return pg_manager.get_process_group("nccl", group) + + +# --------------------------------------------------------------------------- +# Fusion parameters +# --------------------------------------------------------------------------- + + +class FusionOp(IntEnum): + """What post-communication fusion to apply.""" + + NONE = 0 + # all_reduce + residual_add + RMSNorm + RESIDUAL_RMS_NORM = 1 + # reduce_scatter + residual_add + RMSNorm + RS_RESIDUAL_RMS_NORM = 2 + # all_gather + dual RMSNorm (for MLA) + AG_DUAL_RMS_NORM = 3 + + +@dataclass +class FusionParams: + """Optional fusion context passed to fused comm_ops functions. + + Not all fields are used by every ``FusionOp``. Only the relevant + subset is accessed. + """ + + fusion_op: FusionOp = FusionOp.NONE + + # --- For RESIDUAL_RMS_NORM / RS_RESIDUAL_RMS_NORM --- + residual: torch.Tensor | None = None + norm_weight: torch.Tensor | None = None + eps: float = 1e-6 + + # --- For AG_DUAL_RMS_NORM --- + norm_weight_2: torch.Tensor | None = None + eps_2: float = 1e-6 + + # --- For reduce-scatter fusion --- + add_in: torch.Tensor | None = None + residual_reduce_scattered: bool = False + has_partial_norm_out: bool = False + + # --- Shared by RESIDUAL_RMS_NORM / RS_RESIDUAL_RMS_NORM / AG_DUAL_RMS_NORM --- + max_token_num: int = 0 + + # --- For FP8 block quantization --- + block_quant_fp8: bool = False + + # --- General --- + total_num_tokens: int = 0 + trigger_completion_at_end: bool = False + fp32_acc: bool = False + max_sm_to_use: int | None = None + + +# --------------------------------------------------------------------------- +# Basic primitives +# --------------------------------------------------------------------------- + + +def all_reduce( + tensor: torch.Tensor, + group: Group, + backend: CommBackend | None = None, + op: torch.distributed.ReduceOp = torch.distributed.ReduceOp.SUM, +) -> torch.Tensor: + """All-reduce the tensor across the given communication group.""" + if backend is None: + backend = get_global_backend() + return backend.all_reduce(tensor, group, op=op) + + +def all_gather( + tensor: torch.Tensor, + group: Group, + dim: int = -1, + backend: CommBackend | None = None, +) -> torch.Tensor: + """All-gather the tensor across the given communication group.""" + if backend is None: + backend = get_global_backend() + return backend.all_gather(tensor, group, dim) + + +def all_gather_into_tensor( + output: torch.Tensor, + input: torch.Tensor, + group: Group, + backend: CommBackend | None = None, +) -> None: + """All-gather input into a pre-allocated output buffer.""" + if backend is None: + backend = get_global_backend() + backend.all_gather_into_tensor(output, input, group) + + +def reduce_scatter( + tensor: torch.Tensor, + group: Group, + backend: CommBackend | None = None, +) -> torch.Tensor: + """Reduce-scatter the tensor across the given communication group.""" + if backend is None: + backend = get_global_backend() + return backend.reduce_scatter(tensor, group) + + +def all_to_all_single( + output: torch.Tensor, + input: torch.Tensor, + group: Group, + backend: CommBackend | None = None, +) -> None: + """Even-split all_to_all into a pre-allocated output buffer.""" + if backend is None: + backend = get_global_backend() + backend.all_to_all_single(output, input, group) + + +# --------------------------------------------------------------------------- +# Fused ops (comm + residual + norm) +# --------------------------------------------------------------------------- + + +def fused_all_reduce( + tensor: torch.Tensor, + rank: int, + group: Group, + backend: CommBackend | None = None, + fusion_params: FusionParams | None = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """All-reduce with optional fused residual + RMSNorm.""" + if backend is None: + backend = get_global_backend() + + if fusion_params is None or fusion_params.fusion_op == FusionOp.NONE: + return backend.all_reduce(tensor, group) + + if fusion_params.fusion_op == FusionOp.RESIDUAL_RMS_NORM: + return allreduce_residual_rmsnorm( + input_tensor=tensor, + residual=fusion_params.residual, + weight=fusion_params.norm_weight, + rank=rank, + group=_get_process_group(group), + eps=fusion_params.eps, + fp32_acc=fusion_params.fp32_acc, + block_quant_fp8=fusion_params.block_quant_fp8, + residual_reduce_scattered=fusion_params.residual_reduce_scattered, + has_partial_norm_out=fusion_params.has_partial_norm_out, + trigger_completion_at_end=fusion_params.trigger_completion_at_end, + max_sm_to_use=fusion_params.max_sm_to_use, + launch_with_pdl=pdl_enabled(), + ) + + raise ValueError( + f"Unsupported fusion_op {fusion_params.fusion_op} for fused_all_reduce" + ) + + +def fused_reduce_scatter( + tensor: torch.Tensor, + rank: int, + group: Group, + backend: CommBackend | None = None, + fusion_params: FusionParams | None = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Reduce-scatter with optional fused residual + RMSNorm.""" + if backend is None: + backend = get_global_backend() + + if fusion_params is None or fusion_params.fusion_op == FusionOp.NONE: + return backend.reduce_scatter(tensor, group) + + if fusion_params.fusion_op == FusionOp.RS_RESIDUAL_RMS_NORM: + return reducescatter_residual_rmsnorm( + input_tensor=tensor, + weight=fusion_params.norm_weight, + residual=fusion_params.residual, + eps=fusion_params.eps, + rank=rank, + group=_get_process_group(group), + add_in=fusion_params.add_in, + fp32_acc=fusion_params.fp32_acc, + block_quant_fp8=fusion_params.block_quant_fp8, + max_token_num=fusion_params.max_token_num or tensor.shape[0], + launch_with_pdl=pdl_enabled(), + ) + + raise ValueError( + f"Unsupported fusion_op {fusion_params.fusion_op} for fused_reduce_scatter" + ) + + +def fused_all_gather( + tensor: torch.Tensor, + rank: int, + group: Group, + dim: int = -1, + backend: CommBackend | None = None, + fusion_params: FusionParams | None = None, +) -> torch.Tensor | tuple[torch.Tensor, ...]: + """All-gather with optional fused dual-RMSNorm.""" + if backend is None: + backend = get_global_backend() + + if fusion_params is None or fusion_params.fusion_op == FusionOp.NONE: + return backend.all_gather(tensor, group, dim) + + if fusion_params.fusion_op == FusionOp.AG_DUAL_RMS_NORM: + return allgather_dual_rmsnorm( + qkv=tensor, + weight_q_a=fusion_params.norm_weight, + eps_q=fusion_params.eps, + weight_kv_a=fusion_params.norm_weight_2, + eps_kv=fusion_params.eps_2, + rank=rank, + group=_get_process_group(group), + total_num_tokens=fusion_params.total_num_tokens, + max_token_num=fusion_params.max_token_num + or max(tensor.shape[0], fusion_params.total_num_tokens), + fp32_acc=fusion_params.fp32_acc, + block_quant_fp8=fusion_params.block_quant_fp8, + launch_with_pdl=pdl_enabled(), + ) + + raise ValueError( + f"Unsupported fusion_op {fusion_params.fusion_op} for fused_all_gather" + ) + + +# --------------------------------------------------------------------------- +# Token-aware ops (uneven token distribution via TritonRSAG) +# --------------------------------------------------------------------------- + + +def token_all_gather( + tensor: torch.Tensor, + group: Group, + scattered_num_tokens: list[int], + backend=None, +) -> torch.Tensor: + """All-gather with token-aware distribution (TritonRSAG). + + Args: + scattered_num_tokens: Number of tokens on each rank in the group, + e.g. [50, 50, 51, 49] for 4 ranks with 200 total tokens. + """ + if backend is None: + backend = get_global_backend() + return backend.token_all_gather(tensor, group, scattered_num_tokens) + + +def token_reduce_scatter( + tensor: torch.Tensor, + group: Group, + scattered_num_tokens: list[int], + backend=None, +) -> torch.Tensor: + """Reduce-scatter with token-aware distribution (TritonRSAG). + + Args: + scattered_num_tokens: Number of tokens on each rank in the group, + e.g. [50, 50, 51, 49] for 4 ranks with 200 total tokens. + """ + if backend is None: + backend = get_global_backend() + return backend.token_reduce_scatter(tensor, group, scattered_num_tokens) diff --git a/python/tokenspeed/runtime/distributed/device_communicators/custom_all_reduce.py b/python/tokenspeed/runtime/distributed/device_communicators/custom_all_reduce.py new file mode 100755 index 0000000..a6a3b9f --- /dev/null +++ b/python/tokenspeed/runtime/distributed/device_communicators/custom_all_reduce.py @@ -0,0 +1,331 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + + +import ctypes +from contextlib import contextmanager + +import tokenspeed_kernel.ops.communication.flashinfer as custom_ar_ops +import torch +import torch.distributed as dist +from tokenspeed_kernel.platform import current_platform +from tokenspeed_kernel.thirdparty.cuda.cuda_ipc import CudaRTLibrary +from torch.distributed import ProcessGroup + +from tokenspeed.runtime.distributed.device_communicators.utils import ( + in_the_same_node_as, +) +from tokenspeed.runtime.utils import get_colorful_logger +from tokenspeed.runtime.utils.env import envs, global_server_args_dict + +logger = get_colorful_logger(__name__) + +custom_ar = current_platform().is_nvidia + + +def _can_p2p(rank: int, world_size: int) -> bool: + if not global_server_args_dict.get("enable_p2p_check", False): + logger.info("Skipping P2P check and trusting the driver's P2P report.") + return True + + for i in range(world_size): + if i == rank: + continue + return torch.cuda.can_device_access_peer(rank, i) + return True + + +class CustomAllreduce: + _SUPPORTED_WORLD_SIZES = [2, 4, 6, 8] + _MAX_CAR_SIZE = 8192 * 1024 + + # max_size: max supported allreduce size + def __init__( + self, + group: ProcessGroup, + device: 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 + + if dist.get_backend(group) == dist.Backend.NCCL: + raise ValueError("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 + if not isinstance(device, torch.device): + raise TypeError( + f"device must be a torch.device, got {type(device).__name__}" + ) + self.device = device + + # 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 + interconnect = current_platform().interconnect + full_nvlink = ( + interconnect is not None and interconnect.topology == "nvlink_full" + ) + if not full_nvlink and envs.TOKENSPEED_FORCE_FAKE_FULL_NVLINK.get(): + full_nvlink = True + + 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 _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 + + # 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( + custom_ar_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 = custom_ar_ops.init_custom_ar( + self.meta_ptrs, self.rank_data, rank, self.full_nvlink + ) + custom_ar_ops.register_buffer(self._ptr, self.buffer_ptrs) + self.disabled = False + + @staticmethod + def create_shared_buffer( + size_in_bytes: int, group: ProcessGroup | None = 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: ProcessGroup | None = 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 register_graph_buffers(self): + handle, offset = custom_ar_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 + custom_ar_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 + + 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() + ) + + 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 self.world_size == 2 or self.full_nvlink: + return inp_size < self.max_size + 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) + custom_ar_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) + custom_ar_ops.all_reduce_unreg(self._ptr, inp, self.buffer_ptrs[self.rank], 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: + custom_ar_ops.all_reduce(self._ptr, inp, out, 0, 0) + else: + custom_ar_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) -> torch.Tensor | None: + """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(): + 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: + return self.all_reduce(input, registered=False) + + def close(self): + if not self.disabled and self._ptr: + custom_ar_ops.dispose(self._ptr) + 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/tokenspeed/runtime/distributed/device_communicators/pynccl.py b/python/tokenspeed/runtime/distributed/device_communicators/pynccl.py new file mode 100755 index 0000000..28aebbc --- /dev/null +++ b/python/tokenspeed/runtime/distributed/device_communicators/pynccl.py @@ -0,0 +1,299 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + + +from contextlib import contextmanager + +# ===================== import region ===================== +import torch +import torch.distributed as dist +from tokenspeed_kernel.ops.communication.nccl import ( + NCCLLibrary, + buffer_type, + cudaStream_t, + ncclComm_t, + ncclDataTypeEnum, + ncclRedOpTypeEnum, + ncclUniqueId, +) +from torch.distributed import ProcessGroup, ReduceOp + +from tokenspeed.runtime.distributed.utils import StatelessProcessGroup +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +class PyNcclCommunicator: + + def __init__( + self, + group: ProcessGroup | StatelessProcessGroup, + device: int | str | torch.device, + library_path: str | None = 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): + if not dist.is_initialized(): + raise RuntimeError("torch.distributed must be initialized") + if dist.get_backend(group) == dist.Backend.NCCL: + raise ValueError( + "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 + + logger.info("Epsilon 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 + if not isinstance(device, torch.device): + raise TypeError( + f"device must be a torch.device, got {type(device).__name__}" + ) + 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 _check_device(self, tensor: torch.Tensor) -> None: + if tensor.device != self.device: + raise ValueError( + f"this nccl communicator is created to work on {self.device}, " + f"but the input tensor is on {tensor.device}" + ) + + 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" + self._check_device(tensor) + 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 + ): + 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" + self._check_device(input_tensor) + if stream is None: + stream = self.stream + 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, + ): + 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" + self._check_device(input_tensor) + if stream is None: + stream = self.stream + 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 + self._check_device(tensor) + 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 + self._check_device(tensor) + 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 + self._check_device(tensor) + 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), + ) + + @contextmanager + def change_state( + self, enable: bool | None = None, stream: torch.cuda.Stream | None = 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/tokenspeed/runtime/distributed/device_communicators/utils.py b/python/tokenspeed/runtime/distributed/device_communicators/utils.py new file mode 100644 index 0000000..3ae757c --- /dev/null +++ b/python/tokenspeed/runtime/distributed/device_communicators/utils.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Distributed process-group helper utilities.""" + +import contextlib +from multiprocessing import shared_memory +from unittest.mock import patch + +import torch +import torch.distributed +from torch.distributed import ProcessGroup + +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +def _ignore_resource_tracker_register(*args, **kwargs) -> None: + return None + + +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). + """ + if torch.distributed.get_backend(pg) == torch.distributed.Backend.NCCL: + raise ValueError("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", + _ignore_resource_tracker_register, + ): + shm = shared_memory.SharedMemory(name=name) + if shm.buf[: len(magic_message)] == magic_message: + is_in_the_same_node[rank] = 1 + except Exception as exc: + logger.error("Error ignored in is_in_the_same_node: %s", exc) + 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()] diff --git a/python/tokenspeed/runtime/distributed/dp_sampling_comm.py b/python/tokenspeed/runtime/distributed/dp_sampling_comm.py new file mode 100644 index 0000000..ef061d8 --- /dev/null +++ b/python/tokenspeed/runtime/distributed/dp_sampling_comm.py @@ -0,0 +1,446 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Communication helper for Batch-DP speculative verify. + +Here N is num_tokens_per_req, V is the LM-head padded communication +vocab_size, V/TP is the local vocab shard, and reqs_per_rank=pad_bs/TP. +Callers trim swapped logits back to the model config vocab size before +sampling observes token ids. + +swap_batch_vocab maps each rank's full-batch vocab shard +[pad_bs * N, V/TP] to its request shard with full vocab [reqs_per_rank * N, V]. + +gather_verify_outputs maps per-rank verify outputs +predict_local[reqs_per_rank, N], accept_index_local[reqs_per_rank, N], and +accept_length_local[reqs_per_rank] to persistent full-batch buffers +predict_full[pad_bs, N], accept_index_full[pad_bs, N], and +accept_length_full[pad_bs]. +""" + +from __future__ import annotations + +from typing import Literal + +import torch + +from tokenspeed.runtime.distributed.comm_backend import ( + CommBackend, + Group, + get_global_backend, +) +from tokenspeed.runtime.distributed.comm_ops import all_gather_into_tensor +from tokenspeed.runtime.distributed.dp_sampling_swap import ( + swap_batch_vocab as _swap_batch_vocab_nccl, +) +from tokenspeed.runtime.distributed.process_group_manager import ( + process_group_manager as pg_manager, +) +from tokenspeed.runtime.utils import get_colorful_logger +from tokenspeed.runtime.utils.env import envs + +try: + from tokenspeed_kernel.ops.communication.triton import ( + create_dp_sampling_state, + dp_sampling_gather, + dp_sampling_swap, + ) + from tokenspeed_kernel.platform import current_platform + from torch.distributed import _symmetric_memory +except Exception: + create_dp_sampling_state = None + current_platform = None + dp_sampling_gather = None + dp_sampling_swap = None + _symmetric_memory = None + +logger = get_colorful_logger(__name__) + +DpSamplingBackend = Literal["auto", "nccl", "onesided"] +_ResolvedBackend = Literal["nccl", "onesided"] + +ENV_VAR = "TOKENSPEED_DP_SAMPLING_BACKEND" + + +def _env_override() -> DpSamplingBackend | None: + val = envs.TOKENSPEED_DP_SAMPLING_BACKEND.get() + if val in ("auto", "nccl", "onesided"): + return val # type: ignore[return-value] + if val is not None: + raise ValueError(f"{ENV_VAR}={val!r} must be one of 'auto'|'nccl'|'onesided'") + return None + + +def _onesided_available(group: Group) -> bool: + if len(group) <= 1: + return False + if ( + create_dp_sampling_state is None + or current_platform is None + or _symmetric_memory is None + ): + return False + try: + if not current_platform().is_nvidia: + return False + + major, minor = torch.__version__.split("+", 1)[0].split(".")[:2] + if (int(major), int(minor)) < (2, 10): + return False + + return True + except Exception: + return False + + +def _resolve_backend(requested: DpSamplingBackend, group: Group) -> _ResolvedBackend: + env = _env_override() + requested_via_env = env is not None + if env is not None: + requested = env + + if requested == "nccl": + return "nccl" + if requested == "onesided": + if not _onesided_available(group): + fallback_msg = ( + f"Set {ENV_VAR}=nccl or unset {ENV_VAR} to use auto fallback." + if requested_via_env + else f"Set {ENV_VAR}=nccl or use backend='auto' to fall back." + ) + raise RuntimeError( + f"Batch-DP sampling backend='onesided' requested but the one-sided " + f"NVLink kernel is not available for group {group}. " + f"{fallback_msg}" + ) + return "onesided" + + return "onesided" if _onesided_available(group) else "nccl" + + +class DpSamplingComm: + + def __init__( + self, + *, + tp_size: int, + rank: int, + group: Group, + max_pad_bs: int, + num_tokens_per_req: int, + vocab_size: int, + logits_dtype: torch.dtype | None, + backend: DpSamplingBackend = "auto", + fallback_comm_backend: CommBackend | None = None, + device: torch.device | str | None = None, + ): + if tp_size < 1: + raise ValueError(f"tp_size={tp_size}") + if len(group) != tp_size: + raise ValueError( + f"group {group} has {len(group)} ranks but tp_size={tp_size}" + ) + if max_pad_bs % tp_size != 0: + raise ValueError( + f"max_pad_bs={max_pad_bs} must be divisible by tp_size={tp_size}" + ) + if vocab_size % tp_size != 0: + raise ValueError( + f"vocab_size={vocab_size} must be divisible by tp_size={tp_size}" + ) + if num_tokens_per_req < 1: + raise ValueError(f"num_tokens_per_req={num_tokens_per_req}") + + self._tp_size = tp_size + self._rank = rank + self._group = group + self._max_pad_bs = max_pad_bs + self._max_reqs_per_rank = max_pad_bs // tp_size + self._num_tokens_per_req = num_tokens_per_req + self._vocab_size = vocab_size + self._logits_dtype = logits_dtype + self._fallback_backend = fallback_comm_backend or get_global_backend() + self._device = ( + torch.device(device) + if device is not None + else torch.device(f"cuda:{torch.cuda.current_device()}") + ) + + self._backend: _ResolvedBackend = _resolve_backend(backend, group) + self._state = None + + logger.info( + "DpSamplingComm backend=%s tp_size=%d rank=%d max_pad_bs=%d " + "num_tokens_per_req=%d vocab_size=%d", + self._backend, + tp_size, + rank, + max_pad_bs, + num_tokens_per_req, + vocab_size, + ) + + n = num_tokens_per_req + self._predict_full = torch.empty( + max_pad_bs, n, dtype=torch.int32, device=self._device + ) + self._accept_index_full = torch.empty( + max_pad_bs, n, dtype=torch.int32, device=self._device + ) + self._accept_length_full = torch.empty( + max_pad_bs, dtype=torch.int32, device=self._device + ) + self._logprobs_full = torch.empty( + max_pad_bs, n, dtype=torch.float32, device=self._device + ) + + if self._backend == "nccl": + self._combined_local_nccl: torch.Tensor | None = torch.empty( + self._max_reqs_per_rank, + 2 * n + 1, + dtype=torch.int32, + device=self._device, + ) + self._combined_full_nccl: torch.Tensor | None = torch.empty( + max_pad_bs, + 2 * n + 1, + dtype=torch.int32, + device=self._device, + ) + else: + self._combined_local_nccl = None + self._combined_full_nccl = None + + if self._backend == "onesided" and self._logits_dtype is not None: + self._init_onesided() + + @property + def backend(self) -> _ResolvedBackend: + return self._backend + + @property + def fast_path_enabled(self) -> bool: + return self._backend == "onesided" + + @property + def max_pad_bs(self) -> int: + return self._max_pad_bs + + @property + def is_initialized(self) -> bool: + return self._state is not None + + @staticmethod + def _check_shape( + name: str, tensor: torch.Tensor, expected: tuple[int, ...] + ) -> None: + if tuple(tensor.shape) != expected: + raise ValueError(f"{name} shape {tuple(tensor.shape)} != {expected}") + + @staticmethod + def _check_dtype(name: str, tensor: torch.Tensor, expected: torch.dtype) -> None: + if tensor.dtype != expected: + raise TypeError(f"{name} dtype {tensor.dtype} != {expected}") + + def _check_pad_bs(self, pad_bs: int) -> None: + if pad_bs > self._max_pad_bs: + raise ValueError( + f"pad_bs={pad_bs} exceeds max_pad_bs={self._max_pad_bs} " + "(set at construction time)" + ) + if pad_bs % self._tp_size != 0: + raise ValueError( + f"pad_bs={pad_bs} must be divisible by tp_size={self._tp_size}" + ) + + def prepare_verify_outputs(self, logits_dtype: torch.dtype) -> None: + """Initialize one-sided state for verify-only DP sampling routes.""" + if self._backend == "onesided": + if self._state is not None: + return + self._ensure_onesided_state(logits_dtype) + + def swap_batch_vocab( + self, + local_logits: torch.Tensor, + *, + pad_bs: int, + ) -> torch.Tensor: + """Move from vocab shards to request shards. + + Input on each rank is local_logits[pad_bs * N, V_local], where + N=num_tokens_per_req and V_local=V/TP. Output is + [reqs_per_rank * N, V] for this rank's reqs_per_rank=pad_bs/TP + requests. + Returned row local_req * N + d is global request + rank * reqs_per_rank + local_req at draft position d. + """ + self._check_pad_bs(pad_bs) + + if self._backend == "onesided": + self._ensure_onesided_state(local_logits.dtype) + return self._swap_batch_vocab_onesided(local_logits, pad_bs=pad_bs) + + return _swap_batch_vocab_nccl( + local_logits, + tp_size=self._tp_size, + pad_bs=pad_bs, + num_tokens_per_req=self._num_tokens_per_req, + vocab_size=self._vocab_size, + group=self._group, + backend=self._fallback_backend, + ) + + def gather_verify_outputs( + self, + predict_local: torch.Tensor, + accept_index_local: torch.Tensor, + accept_length_local: torch.Tensor, + *, + pad_bs: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Gather local verify outputs into full padded-batch outputs. + + Inputs are predict_local[reqs_per_rank, N], + accept_index_local[reqs_per_rank, N], and + accept_length_local[reqs_per_rank]. + Returns predict_full[pad_bs, N], accept_index_full[pad_bs, N], and + accept_length_full[pad_bs]. + Row r from source rank src lands at src * reqs_per_rank + r. + Callers slice the real [0:bs] prefix and ignore phantom rows. + """ + self._check_pad_bs(pad_bs) + reqs_per_rank = pad_bs // self._tp_size + n = self._num_tokens_per_req + + self._check_shape("predict_local", predict_local, (reqs_per_rank, n)) + self._check_shape("accept_index_local", accept_index_local, (reqs_per_rank, n)) + self._check_shape("accept_length_local", accept_length_local, (reqs_per_rank,)) + self._check_dtype("predict_local", predict_local, torch.int32) + self._check_dtype("accept_index_local", accept_index_local, torch.int32) + self._check_dtype("accept_length_local", accept_length_local, torch.int32) + + if self._backend == "onesided": + return self._gather_verify_outputs_onesided( + predict_local, + accept_index_local, + accept_length_local, + pad_bs=pad_bs, + ) + + if self._combined_local_nccl is None or self._combined_full_nccl is None: + raise RuntimeError("NCCL DP sampling buffers are not initialized") + combined_local = self._combined_local_nccl[:reqs_per_rank] + combined_local[:, :n].copy_(predict_local) + combined_local[:, n : 2 * n].copy_(accept_index_local) + combined_local[:, 2 * n].copy_(accept_length_local) + + combined_full = self._combined_full_nccl[:pad_bs] + all_gather_into_tensor( + combined_full, + combined_local, + self._group, + backend=self._fallback_backend, + ) + + predict_full = self._predict_full[:pad_bs] + accept_index_full = self._accept_index_full[:pad_bs] + accept_length_full = self._accept_length_full[:pad_bs] + predict_full.copy_(combined_full[:, :n]) + accept_index_full.copy_(combined_full[:, n : 2 * n]) + accept_length_full.copy_(combined_full[:, 2 * n]) + return predict_full, accept_index_full, accept_length_full + + def gather_verify_logprobs( + self, + logprobs_local: torch.Tensor, + *, + pad_bs: int, + ) -> torch.Tensor: + """Gather per-token scalar logprobs into full padded-batch order.""" + self._check_pad_bs(pad_bs) + reqs_per_rank = pad_bs // self._tp_size + n = self._num_tokens_per_req + self._check_shape("logprobs_local", logprobs_local, (reqs_per_rank, n)) + logprobs_full = self._logprobs_full[:pad_bs] + all_gather_into_tensor( + logprobs_full, + logprobs_local.contiguous(), + self._group, + backend=self._fallback_backend, + ) + return logprobs_full + + def _init_onesided(self) -> None: + if self._logits_dtype is None: + raise RuntimeError("DP sampling logits dtype is not initialized") + if create_dp_sampling_state is None: + raise RuntimeError("one-sided DP sampling state creation is unavailable") + + self._state = create_dp_sampling_state( + group=pg_manager.get_process_group("nccl", self._group), + rank_in_group=self._rank, + tp_size=self._tp_size, + max_pad_bs=self._max_pad_bs, + num_tokens_per_req=self._num_tokens_per_req, + vocab_size=self._vocab_size, + logits_dtype=self._logits_dtype, + device=self._device, + ) + + def _ensure_onesided_state(self, logits_dtype: torch.dtype) -> None: + if self._state is not None: + if self._logits_dtype != logits_dtype: + raise RuntimeError( + f"DP sampling logits dtype changed from {self._logits_dtype} " + f"to {logits_dtype}" + ) + return + self._logits_dtype = logits_dtype + self._init_onesided() + + def _swap_batch_vocab_onesided( + self, local_logits: torch.Tensor, *, pad_bs: int + ) -> torch.Tensor: + if self._state is None: + raise RuntimeError("one-sided DP sampling state is not initialized") + if dp_sampling_swap is None: + raise RuntimeError("one-sided DP sampling swap is unavailable") + return dp_sampling_swap(self._state, local_logits, pad_bs=pad_bs) + + def _gather_verify_outputs_onesided( + self, + predict_local: torch.Tensor, + accept_index_local: torch.Tensor, + accept_length_local: torch.Tensor, + *, + pad_bs: int, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if self._state is None: + raise RuntimeError("one-sided DP sampling state is not initialized") + if dp_sampling_gather is None: + raise RuntimeError("one-sided DP sampling gather is unavailable") + return dp_sampling_gather( + self._state, + predict_local, + accept_index_local, + accept_length_local, + pad_bs=pad_bs, + ) diff --git a/python/tokenspeed/runtime/distributed/dp_sampling_swap.py b/python/tokenspeed/runtime/distributed/dp_sampling_swap.py new file mode 100644 index 0000000..69ddf62 --- /dev/null +++ b/python/tokenspeed/runtime/distributed/dp_sampling_swap.py @@ -0,0 +1,77 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""NCCL fallback for Batch-DP logits shape swap.""" + +from __future__ import annotations + +import torch + +from tokenspeed.runtime.distributed.comm_backend import CommBackend, Group +from tokenspeed.runtime.distributed.comm_ops import all_to_all_single + + +def swap_batch_vocab( + local_logits: torch.Tensor, + *, + tp_size: int, + pad_bs: int, + num_tokens_per_req: int, + vocab_size: int, + group: Group, + backend: CommBackend | None = None, +) -> torch.Tensor: + """Move logits from vocab shards to request shards. + + Each rank starts with local_logits[pad_bs * N, V_local] for the full + padded batch and its local vocab slice, where V_local=V/TP. The result is + [reqs_per_rank * N, V] for this rank's reqs_per_rank=pad_bs/TP requests. + Returned row local_req * N + d is global request + rank * reqs_per_rank + local_req at draft position d. + """ + if pad_bs % tp_size != 0: + raise ValueError( + f"swap_batch_vocab: pad_bs={pad_bs} must be divisible by tp_size={tp_size}" + ) + if vocab_size % tp_size != 0: + raise ValueError( + f"swap_batch_vocab: vocab_size={vocab_size} must be divisible by tp_size={tp_size}" + ) + + reqs_per_rank = pad_bs // tp_size + v_local = vocab_size // tp_size + n = num_tokens_per_req + + expected_shape = (pad_bs * n, v_local) + if tuple(local_logits.shape) != expected_shape: + raise ValueError( + f"swap_batch_vocab: local_logits shape {tuple(local_logits.shape)} " + f"!= expected {expected_shape} (pad_bs={pad_bs}, N={n}, V/TP={v_local})" + ) + + recv = torch.empty_like(local_logits) + all_to_all_single(recv, local_logits, group, backend=backend) + + return ( + recv.view(tp_size, reqs_per_rank, n, v_local) + .permute(1, 2, 0, 3) + .contiguous() + .view(reqs_per_rank * n, vocab_size) + ) diff --git a/python/tokenspeed/runtime/distributed/mapping.py b/python/tokenspeed/runtime/distributed/mapping.py new file mode 100644 index 0000000..f924f21 --- /dev/null +++ b/python/tokenspeed/runtime/distributed/mapping.py @@ -0,0 +1,392 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import math +from functools import cached_property + +Group = tuple[int, ...] + + +def _resolve_parallelism_sizes(world_size: int, *sizes: int | None) -> tuple[int, ...]: + """Resolve the parallelism sizes given world_size. + + `sizes` is ordered innermost (fastest-varying) to outermost. + """ + assert all(x is None or x > 0 for x in sizes) + + resolved = [x for x in sizes] + num_to_resolve = sum(x is None for x in sizes) + if num_to_resolve > 0: + provided_size = math.prod(x for x in sizes if x is not None) + assert provided_size <= world_size + assert world_size % provided_size == 0 + resolved_size = world_size // provided_size + + for index, size in enumerate(resolved): + if size is None: + resolved[index] = resolved_size + resolved_size = 1 + + assert math.prod(resolved) == world_size + return tuple(resolved) + + +def _make_parallelism_rank(rank: int, size: int, stride: int = 1) -> int: + """Return the rank of given size and stride.""" + return (rank // stride) % size + + +def _make_parallelism_group(rank: int, size: int, stride: int = 1) -> Group: + """Return the group of ranks of given size and stride.""" + base = rank - (rank // stride % size) * stride + return tuple(base + j * stride for j in range(size)) + + +class MappingBase: + + def __init__(self, rank: int | None = None, world_size: int = 1): + assert rank is None or rank >= 0 + self._rank = rank + assert world_size > 0 + self._world_size = world_size + + @property + def rank(self) -> int: + assert self._rank is not None, "rank is not initialized" + return self._rank + + @rank.setter + def rank(self, rank: int): + assert self._rank is None, "rank is already initialized" + assert rank >= 0 + self._rank = rank + self._on_rank_initialized(rank) + + def _on_rank_initialized(self, rank: int): + return None + + @property + def world_size(self) -> int: + return self._world_size + + @cached_property + def world_group(self) -> Group: + return _make_parallelism_group(self.rank, self.world_size, stride=1) + + +class DenseLayerMapping(MappingBase): + + def __init__( + self, + rank: int | None = None, + world_size: int = 1, + tp_size: int | None = None, + dp_size: int | None = None, + ): + super().__init__(rank, world_size) + self.tp_size, self.dp_size = _resolve_parallelism_sizes( + self.world_size, tp_size, dp_size + ) + + @cached_property + def has_tp(self) -> bool: + return self.tp_size > 1 + + @cached_property + def tp_rank(self) -> int: + return _make_parallelism_rank(self.rank, self.tp_size, stride=1) + + @cached_property + def tp_group(self) -> Group: + return _make_parallelism_group(self.rank, self.tp_size, stride=1) + + @cached_property + def has_dp(self) -> bool: + return self.dp_size > 1 + + @cached_property + def dp_rank(self) -> int: + return _make_parallelism_rank(self.rank, self.dp_size, stride=self.tp_size) + + @cached_property + def dp_group(self) -> Group: + return _make_parallelism_group(self.rank, self.dp_size, stride=self.tp_size) + + +class AttentionLayerMapping(MappingBase): + + def __init__( + self, + rank: int | None = None, + world_size: int = 1, + tp_size: int | None = None, + cp_size: int | None = None, + dp_size: int | None = None, + ): + super().__init__(rank, world_size) + self.tp_size, self.cp_size, self.dp_size = _resolve_parallelism_sizes( + self.world_size, tp_size, cp_size, dp_size + ) + + @cached_property + def has_tp(self) -> bool: + return self.tp_size > 1 + + @cached_property + def tp_rank(self) -> int: + return _make_parallelism_rank(self.rank, self.tp_size, stride=1) + + @cached_property + def tp_group(self) -> Group: + return _make_parallelism_group(self.rank, self.tp_size, stride=1) + + @cached_property + def has_cp(self) -> bool: + return self.cp_size > 1 + + @cached_property + def cp_rank(self) -> int: + return _make_parallelism_rank(self.rank, self.cp_size, stride=self.tp_size) + + @cached_property + def cp_group(self) -> Group: + return _make_parallelism_group(self.rank, self.cp_size, stride=self.tp_size) + + @cached_property + def has_dp(self) -> bool: + return self.dp_size > 1 + + @cached_property + def dp_rank(self) -> int: + return _make_parallelism_rank( + self.rank, self.dp_size, stride=self.tp_size * self.cp_size + ) + + @cached_property + def dp_group(self) -> Group: + return _make_parallelism_group( + self.rank, self.dp_size, stride=self.tp_size * self.cp_size + ) + + def scatter_index(self, rank: int) -> int: + """Index of ``rank`` in a dp-major/tp-minor scattered token count + table; cp peers share their dp group's tp split.""" + tp_rank = _make_parallelism_rank(rank, self.tp_size, stride=1) + dp_rank = _make_parallelism_rank( + rank, self.dp_size, stride=self.tp_size * self.cp_size + ) + return dp_rank * self.tp_size + tp_rank + + +class MoeLayerMapping(MappingBase): + def __init__( + self, + rank: int | None = None, + world_size: int = 1, + tp_size: int | None = None, + ep_size: int | None = None, + dp_size: int | None = None, + ): + super().__init__(rank, world_size) + self.tp_size, self.ep_size, self.dp_size = _resolve_parallelism_sizes( + self.world_size, tp_size, ep_size, dp_size + ) + + @cached_property + def has_tp(self) -> bool: + return self.tp_size > 1 + + @cached_property + def tp_rank(self) -> int: + return _make_parallelism_rank(self.rank, self.tp_size, stride=1) + + @cached_property + def tp_group(self) -> Group: + return _make_parallelism_group(self.rank, self.tp_size, stride=1) + + @cached_property + def has_ep(self) -> bool: + return self.ep_size > 1 + + @cached_property + def ep_rank(self) -> int: + return _make_parallelism_rank(self.rank, self.ep_size, stride=self.tp_size) + + @cached_property + def ep_group(self) -> Group: + return _make_parallelism_group(self.rank, self.ep_size, stride=self.tp_size) + + @cached_property + def has_tp_ep(self) -> bool: + return self.tp_ep_size > 1 + + @cached_property + def tp_ep_size(self) -> int: + return self.tp_size * self.ep_size + + @cached_property + def tp_ep_rank(self) -> int: + return _make_parallelism_rank(self.rank, self.tp_ep_size, stride=1) + + @cached_property + def tp_ep_group(self) -> Group: + return _make_parallelism_group(self.rank, self.tp_ep_size, stride=1) + + @cached_property + def has_dp(self) -> bool: + return self.dp_size > 1 + + @cached_property + def dp_rank(self) -> int: + return _make_parallelism_rank( + self.rank, self.dp_size, stride=self.tp_size * self.ep_size + ) + + @cached_property + def dp_group(self) -> Group: + return _make_parallelism_group( + self.rank, self.dp_size, stride=self.tp_size * self.ep_size + ) + + +class VisionTowerMapping(MappingBase): + """Parallel mapping for vision encoders. Vision layers run colocated and + share the attention TP group; non-colocated deployments should run the + encoder out-of-engine (EPD-style workers + gateway dispatch). + """ + + def __init__( + self, + rank: int | None = None, + world_size: int = 1, + tp_size: int | None = None, + ): + super().__init__(rank, world_size) + (self.tp_size,) = _resolve_parallelism_sizes(self.world_size, tp_size) + + @cached_property + def has_tp(self) -> bool: + return self.tp_size > 1 + + @cached_property + def tp_rank(self) -> int: + return _make_parallelism_rank(self.rank, self.tp_size, stride=1) + + @cached_property + def tp_group(self) -> Group: + return _make_parallelism_group(self.rank, self.tp_size, stride=1) + + +class Mapping(MappingBase): + + def __init__( + self, + rank: int | None = None, + world_size: int = 1, + *, + attn_tp_size: int | None = None, + attn_cp_size: int | None = None, + attn_dp_size: int | None = None, + dense_tp_size: int | None = None, + dense_dp_size: int | None = None, + moe_tp_size: int | None = None, + moe_ep_size: int | None = None, + moe_dp_size: int | None = None, + nprocs_per_node: int | None = None, + nnodes: int | None = None, + base_gpu_id: int = 0, + gpu_id_step: int = 1, + ): + super().__init__(rank, world_size) + self.attn = AttentionLayerMapping( + rank=rank, + world_size=world_size, + tp_size=attn_tp_size, + cp_size=attn_cp_size, + dp_size=attn_dp_size, + ) + self.dense = DenseLayerMapping( + rank=rank, + world_size=world_size, + tp_size=dense_tp_size, + dp_size=dense_dp_size, + ) + self.moe = MoeLayerMapping( + rank=rank, + world_size=world_size, + tp_size=moe_tp_size, + ep_size=moe_ep_size, + dp_size=moe_dp_size, + ) + # Vision tower runs colocated on the attention TP group. + self.vision = VisionTowerMapping( + rank=rank, + world_size=self.attn.tp_size, + tp_size=self.attn.tp_size, + ) + self.nprocs_per_node, self.nnodes = _resolve_parallelism_sizes( + self.world_size, nprocs_per_node, nnodes + ) + assert base_gpu_id >= 0 + assert gpu_id_step > 0 + self.base_gpu_id = base_gpu_id + self.gpu_id_step = gpu_id_step + + def _on_rank_initialized(self, rank: int): + self.attn.rank = rank + self.dense.rank = rank + self.moe.rank = rank + self.vision.rank = rank + + @cached_property + def has_attn_tp(self) -> bool: + return self.attn.has_tp + + @cached_property + def has_attn_cp(self) -> bool: + return self.attn.has_cp + + @cached_property + def has_attn_dp(self) -> bool: + return self.attn.has_dp + + @cached_property + def node_rank(self) -> int: + return self.rank // self.nprocs_per_node + + @cached_property + def local_rank(self) -> int: + return self.rank % self.nprocs_per_node + + @cached_property + def gpu_id(self) -> int: + return self.base_gpu_id + self.local_rank * self.gpu_id_step + + def __repr__(self) -> str: + rank_str = str(self._rank) if self._rank is not None else "?" + lines = [ + f"Mapping(rank={rank_str}, world_size={self.world_size})", + f" Cluster : {self.nnodes} node(s) x {self.nprocs_per_node} proc(s)", + f" Attention: tp={self.attn.tp_size} cp={self.attn.cp_size} dp={self.attn.dp_size}", + f" Vision: tp={self.vision.tp_size}", + f" Dense : tp={self.dense.tp_size} dp={self.dense.dp_size}", + f" MoE : tp={self.moe.tp_size} ep={self.moe.ep_size} dp={self.moe.dp_size}", + ] + return "\n".join(lines) diff --git a/python/tokenspeed/runtime/distributed/process_group_manager.py b/python/tokenspeed/runtime/distributed/process_group_manager.py new file mode 100644 index 0000000..0bd35f9 --- /dev/null +++ b/python/tokenspeed/runtime/distributed/process_group_manager.py @@ -0,0 +1,110 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Helpers for initializing and caching torch distributed process groups.""" + +from datetime import timedelta + +import torch.distributed as dist + +from tokenspeed.runtime.distributed.mapping import Group, Mapping + + +def _make_all_groups(group: Group) -> list[Group]: + """Enumerate all groups with the same size and stride pattern as ``group``.""" + size = len(group) + stride = group[1] - group[0] if len(group) > 1 else 1 + block = size * stride + world_size = dist.get_world_size() + + groups = [] + for base in range(0, world_size, block): + for offset in range(stride): + g = tuple(base + offset + i * stride for i in range(size)) + groups.append(g) + return groups + + +class ProcessGroupManager: + def __init__(self): + self._process_groups: dict[str, dict[Group, dist.ProcessGroup]] = {} + + def init_distributed( + self, + mapping: Mapping, + distributed_init_method: str = "env://", + backend: str = "nccl", + timeout: int | None = None, + ) -> None: + if not dist.is_initialized(): + if distributed_init_method is None: + raise ValueError( + "distributed_init_method must be provided when initializing distributed environment" + ) + if timeout is not None: + if not isinstance(timeout, int): + raise TypeError("timeout must be a number") + if timeout <= 0: + raise ValueError("timeout must be positive") + timeout = timedelta(seconds=timeout) + + dist.init_process_group( + backend=backend, + init_method=distributed_init_method, + world_size=mapping.world_size, + rank=mapping.rank, + timeout=timeout, + ) + + def register_process_group( + self, backend: str, group: Group, process_group: dist.ProcessGroup + ) -> None: + if backend not in self._process_groups: + self._process_groups[backend] = {} + self._process_groups[backend][group] = process_group + + def get_process_group(self, backend: str, group: Group): + return self._process_groups[backend][group] + + def has_process_group(self, backend: str, group: Group) -> bool: + if backend not in self._process_groups: + return False + return group in self._process_groups[backend] + + def init_process_group( + self, group: Group, backend: str | list[str] | None = None + ) -> None: + if backend is None: + backends = ["nccl", "gloo"] + elif isinstance(backend, str): + backends = [backend] + else: + backends = backend + + for backend in backends: + if self.has_process_group(backend, group): + continue + for g in _make_all_groups(group): + pg = dist.new_group(g, backend=backend) + if g == group: + self.register_process_group(backend, g, pg) + + +process_group_manager = ProcessGroupManager() diff --git a/python/tokenspeed/runtime/distributed/utils.py b/python/tokenspeed/runtime/distributed/utils.py new file mode 100755 index 0000000..b2c6b8c --- /dev/null +++ b/python/tokenspeed/runtime/distributed/utils.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Utilities for distributed runtime helpers and stateless coordination.""" + +import dataclasses +import pickle +import time +from collections import deque +from collections.abc import Sequence +from typing import Any + +import torch +from torch.distributed import TCPStore + +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +def ensure_divisibility(numerator, denominator) -> None: + """Ensure that numerator is divisible by the denominator.""" + if numerator % denominator != 0: + raise ValueError(f"{numerator} is not divisible by {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) + # 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 + + +@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 + + 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): + if self.rank >= self.world_size: + raise ValueError( + f"rank={self.rank} must be less than world_size={self.world_size}" + ) + self.broadcast_recv_src_counter = {i: 0 for i in range(self.world_size)} + + def expire_data(self) -> None: + """Expire data that is older than `data_expiration_seconds` seconds.""" + while self.entries: + # check the oldest entry + key, timestamp = self.entries[0] + if time.time() - timestamp > self.data_expiration_seconds: + self.store.delete_key(key) + self.entries.popleft() + else: + break + + def broadcast_obj(self, obj: Any | None, 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}/{self.broadcast_send_counter}" + self.store.set(key, pickle.dumps(obj)) + self.broadcast_send_counter += 1 + self.entries.append((key, time.time())) + return obj + else: + key = f"broadcast_from/{src}/{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 barrier(self) -> None: + """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. + """ + 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/tokenspeed/runtime/engine/__init__.py b/python/tokenspeed/runtime/engine/__init__.py new file mode 100644 index 0000000..96422c1 --- /dev/null +++ b/python/tokenspeed/runtime/engine/__init__.py @@ -0,0 +1,21 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Runtime engine package.""" diff --git a/python/tokenspeed/runtime/engine/aio_rwlock.py b/python/tokenspeed/runtime/engine/aio_rwlock.py new file mode 100755 index 0000000..c8f28c7 --- /dev/null +++ b/python/tokenspeed/runtime/engine/aio_rwlock.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the FluentLLM project +# SPDX-FileCopyrightText: Copyright 2023-2024 SGLang Team +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +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/tokenspeed/runtime/engine/async_llm.py b/python/tokenspeed/runtime/engine/async_llm.py new file mode 100755 index 0000000..a5e532c --- /dev/null +++ b/python/tokenspeed/runtime/engine/async_llm.py @@ -0,0 +1,784 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""AsyncLLM is the main-process async frontend. + +Owns request intake, per-request state, scheduler IPC, and the +output-dispatch loop. Inherits from ``EngineClient`` (explicit +structural conformance) and ``SchedulerControlClient`` (scheduler +control-plane helpers). +""" + +import asyncio +import copy +import logging +import os +import signal +import sys +import threading +import time +import uuid +from collections import deque +from collections.abc import Awaitable +from enum import Enum +from http import HTTPStatus +from typing import ( + Any, + Generic, + TypeVar, +) + +import uvloop + +from tokenspeed.runtime.configs.model_config import ModelConfig +from tokenspeed.runtime.engine.aio_rwlock import RWLock +from tokenspeed.runtime.engine.collector import RequestOutputCollector +from tokenspeed.runtime.engine.core_client import EngineCoreClient +from tokenspeed.runtime.engine.exceptions import EngineGenerateError +from tokenspeed.runtime.engine.input_processor import InputProcessor +from tokenspeed.runtime.engine.io_struct import ( + AbortReq, + BatchEmbeddingOut, + BatchStrOut, + BatchTokenIDOut, + CloseSessionReqInput, + ConfigureLoggingReq, + EmbeddingReqInput, + FlushCacheReqInput, + FlushCacheReqOutput, + GenerateReqInput, + GetLoadReqInput, + HealthCheckOutput, + OpenSessionReqInput, + OpenSessionReqOutput, + TokenizedEmbeddingReqInput, + TokenizedGenerateReqInput, + UpdateWeightFromDiskReqInput, + UpdateWeightFromDiskReqOutput, + WatchLoadUpdateReq, +) +from tokenspeed.runtime.engine.output_processor import OutputProcessor, ReqState +from tokenspeed.runtime.engine.parallel_sampling import ( + prepare_parallel_sampling_replica, + prepare_prefix_warmup, +) +from tokenspeed.runtime.engine.protocol import EngineClient +from tokenspeed.runtime.engine.scheduler_control_client import ( + SchedulerControlClient, +) +from tokenspeed.runtime.metrics.collector import RequestMetrics +from tokenspeed.runtime.pd.utils import ( + DisaggregationMode, + KVClassType, + TransferBackend, + get_kv_class, +) +from tokenspeed.runtime.utils import ( + dataclass_to_string_truncated, + get_colorful_logger, +) +from tokenspeed.runtime.utils.dispatch import TypeBasedDispatcher +from tokenspeed.runtime.utils.exceptions import get_exception_traceback +from tokenspeed.runtime.utils.hf_transformers_utils import get_tokenizer +from tokenspeed.runtime.utils.process import kill_process_tree +from tokenspeed.runtime.utils.server_args import PortArgs, ServerArgs + +asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) + +logger = get_colorful_logger(__name__) + + +def _ignore_health_check_output(_: HealthCheckOutput) -> None: + return None + + +class ServerStatus(Enum): + Up = "Up" + Starting = "Starting" + UnHealthy = "UnHealthy" + Crashed = "Crashed" + + +class AsyncLLM(SchedulerControlClient, EngineClient): + """Main-process async frontend for the tokenspeed runtime. + + Owns request intake, per-request state, scheduler IPC, and the + output-dispatch loop. Structurally satisfies :class:`EngineClient` + via the explicit inheritance declaration above. + """ + + def __init__( + self, + server_args: ServerArgs, + port_args: PortArgs, + ): + # Parse args + self.server_args = server_args + self.enable_metrics = server_args.enable_metrics + self.log_requests = server_args.enable_log_requests + self.log_requests_level = server_args.log_requests_level + self.logger = logger + + # Init inter-process communication (scheduler IPC owned by EngineCoreClient). + self.engine_core_client = EngineCoreClient(port_args) + + # Read model args + self.model_path = server_args.model + self.served_model_name = server_args.served_model_name + self.model_config = ModelConfig( + server_args.model, + trust_remote_code=server_args.trust_remote_code, + revision=server_args.revision, + context_length=server_args.max_model_len, + model_override_args=server_args.hf_overrides, + dtype=server_args.dtype, + quantization=server_args.quantization, + server_args=server_args, + ) + + self.is_generation = self.model_config.is_generation + self.is_image_gen = self.model_config.is_image_gen + self.context_len = self.model_config.context_len + self.image_token_id = self.model_config.image_token_id + # Create tokenizer. The engine never preprocesses images -- the SMG + # gateway ships precomputed multimodal inputs -- so even multimodal + # models only need the tokenizer, not the full HF AutoProcessor. + if server_args.skip_tokenizer_init: + self.tokenizer = None + else: + self.tokenizer = get_tokenizer( + server_args.tokenizer, + tokenizer_mode=server_args.tokenizer_mode, + trust_remote_code=server_args.trust_remote_code, + revision=server_args.revision, + architectures=self.model_config.hf_config.architectures, + ) + if self.model_config.is_multimodal: + os.environ["TOKENIZERS_PARALLELISM"] = "false" + # Store states + self.no_create_loop = False + self.rid_to_state: dict[str, ReqState] = {} + self.gracefully_exit = False + self.last_receive_tstamp = 0 + self.dump_requests_folder = "" # By default do not dump + self.dump_requests_threshold = 1000 + self.dump_request_list: list[tuple] = [] + self.log_request_metadata = self.get_log_request_metadata() + self.server_status = ServerStatus.Starting + + # The event to notify the weight sync is finished. + self.model_update_lock = RWLock() + self.model_update_result: Awaitable[UpdateWeightFromDiskReqOutput] | None = None + self.asyncio_tasks = set() + + # For session info + self.session_futures = {} # session_id -> asyncio event + + # Set after scheduler is initialized + self.max_req_input_len = None + + self.metrics = RequestMetrics( + labels={ + "model_name": self.server_args.served_model_name, + "app_key": self.server_args.app_key, + }, + enabled=( + self.enable_metrics + and "prometheus" in (server_args.metrics_reporters or []) + ), + ) + + self.output_processor = OutputProcessor(self) + + self._result_dispatcher = TypeBasedDispatcher( + [ + ( + ( + BatchStrOut, + BatchEmbeddingOut, + BatchTokenIDOut, + ), + self.output_processor.handle_batch_output, + ), + (OpenSessionReqOutput, self._handle_open_session_req_output), + ( + UpdateWeightFromDiskReqOutput, + self._handle_update_weights_from_disk_req_output, + ), + (HealthCheckOutput, _ignore_health_check_output), + ] + ) + + self.disaggregation_mode = DisaggregationMode( + self.server_args.disaggregation_mode + ) + self.transfer_backend = TransferBackend( + self.server_args.disaggregation_transfer_backend + ) + # for disaggregation, start kv bootstrap server on prefill + if self.disaggregation_mode == DisaggregationMode.PREFILL: + # only start bootstrap server on prefill tm + kv_bootstrap_server_class = get_kv_class( + self.transfer_backend, KVClassType.BOOTSTRAP_SERVER + ) + self.bootstrap_server = kv_bootstrap_server_class( + self.server_args.disaggregation_bootstrap_port + ) + + self.init_communicators(server_args) + + # Tokenization lives in :class:`InputProcessor`; see + # :meth:`_tokenize_one_request` for the delegation. + self.input_processor = InputProcessor(self) + + async def generate_request( + self, + obj: GenerateReqInput | EmbeddingReqInput, + ): + created_time = time.time() + + self.auto_create_handle_loop() + + self.input_processor.validate_request(obj) + + obj.normalize_batch_and_arguments() + + if self.log_requests: + max_length, skip_names, _ = self.log_request_metadata + logger.info( + "Receive: obj=%s", + dataclass_to_string_truncated(obj, max_length, skip_names=skip_names), + ) + + async with self.model_update_lock.reader_lock: + is_single = obj.is_single + if is_single: + tokenized_obj = await self._tokenize_one_request(obj) + self._send_one_request(obj, tokenized_obj, created_time) + async for response in self._wait_one_response(obj): + yield response + else: + async for response in self._handle_batch_request(obj, created_time): + yield response + + async def _tokenize_one_request( + self, + obj: GenerateReqInput | EmbeddingReqInput, + ) -> TokenizedGenerateReqInput | TokenizedEmbeddingReqInput: + """Delegate to :class:`InputProcessor`. + + The tokenization body lives in + ``InputProcessor.tokenize_one_request``. If the input-side + surface ever grows (e.g. multimodal routing), it happens in + that module — not here. + """ + return await self.input_processor.tokenize_one_request(obj) + + def _send_one_request( + self, + obj: GenerateReqInput | EmbeddingReqInput, + tokenized_obj: TokenizedGenerateReqInput | TokenizedEmbeddingReqInput, + created_time: float | None = None, + ): + state = ReqState( + RequestOutputCollector(), + False, + asyncio.Event(), + obj, + created_time=created_time, + tokenized_time=tokenized_obj.created_time, + ) + self.rid_to_state[obj.rid] = state + mm_inputs = getattr(tokenized_obj, "multimodal_inputs", None) + if mm_inputs is not None: + mm_inputs.publish_shm_features() + self.engine_core_client.send_to_scheduler.send_pyobj(tokenized_obj) + + def submit_encode(self, encode_request) -> None: + """Send an EPD encode request to the encode-worker scheduler subprocess. + + Fire-and-forget over the same scheduler-input channel the LM uses (the + encode subprocess runs ``run_encode_loop``, not the LM EventLoop). The + encode worker runs the vision tower and ships the embeddings to prefill + over Mooncake; the gateway gets no streamed response, only the gRPC ack. + ``encode_request`` is an ``encode_worker.EncodeRequest`` whose multimodal + tensors are CPU (gateway-reconstructed) and pickle over ZMQ directly; shm + publishing of the pixels is a follow-up optimization. + """ + self.engine_core_client.send_to_scheduler.send_pyobj(encode_request) + + async def _wait_one_response( + self, + obj: GenerateReqInput | EmbeddingReqInput, + ): + """Wait for the response of one request. + + Cancellation contract: callers (FastAPI route handlers, the sync + ``LLM`` bridge, RL-trainer drivers, etc.) signal client disconnect + via ``asyncio.CancelledError`` — not via a polled + ``request.is_disconnected()`` check. If the task driving this + generator is cancelled mid-wait, the ``finally`` below drops the + rid from ``rid_to_state`` and fires an ``AbortReq`` at the + scheduler so no per-request state leaks. + """ + state = self.rid_to_state[obj.rid] + + try: + while True: + await state.event.wait() + + out = state.collector.take() + if out is None: + state.event.clear() + continue + if state.finished: + if self.log_requests: + max_length, skip_names, out_skip_names = ( + self.log_request_metadata + ) + if self.model_config.is_multimodal_gen: + msg = f"Finish: obj={dataclass_to_string_truncated(obj, max_length, skip_names=skip_names)}" + else: + if ( + isinstance(obj, GenerateReqInput) + and obj.input_ids is not None + and obj.text is None + ): + if self.tokenizer is not None: + obj.text = self.tokenizer.decode( + obj.input_ids, + skip_special_tokens=getattr( + obj.sampling_params, + "skip_special_tokens", + False, + ), + ) + else: + obj.text = "" + msg = f"Finish: obj={dataclass_to_string_truncated(obj, max_length, skip_names=skip_names)}, out={dataclass_to_string_truncated(out, max_length, skip_names=out_skip_names)}" + logger.info(msg) + del self.rid_to_state[obj.rid] + + # Check if this was an abort/error created by scheduler + if isinstance(out["meta_info"].get("finish_reason"), dict): + finish_reason = out["meta_info"]["finish_reason"] + if finish_reason.get("type") == "abort": + if ( + finish_reason.get("status_code") + == HTTPStatus.BAD_REQUEST + ): + raise EngineGenerateError(finish_reason["message"]) + + yield out + break + + state.event.clear() + if state.collector.has_pending(): + state.event.set() + + if obj.stream: + yield out + # else: non-stream path falls through and waits for the + # final chunk; external cancellation wakes us via + # ``asyncio.CancelledError`` from ``state.event.wait()``. + finally: + # Idempotent cleanup split on ``state.finished``: + # + # * Normal-finish path: the yield loop above already did + # ``del self.rid_to_state[obj.rid]`` at ``state.finished``. + # We defensively ``pop`` with a default so a second exit + # through ``finally`` (e.g. when the yield itself raised + # after the del) is a no-op. + # + # * Abandoned path (CancelledError / unexpected exception): + # the rid is still in ``rid_to_state``. Call + # ``abort_request`` which **both** removes it from the + # state map **and** sends ``AbortReq`` to the scheduler. + # Ordering matters: ``abort_request`` early-returns if + # the rid is already gone, so we must not pop first. + if state.finished: + self.rid_to_state.pop(obj.rid, None) + else: + self.abort_request(obj.rid) + + async def _handle_batch_request( + self, + obj: GenerateReqInput | EmbeddingReqInput, + created_time: float | None = None, + ): + batch_size = obj.batch_size + + generators = [] + rids = [] + if getattr(obj, "parallel_sample_num", 1) == 1: + # Send all requests + for i in range(batch_size): + tmp_obj = obj[i] + tokenized_obj = await self._tokenize_one_request(tmp_obj) + self._send_one_request(tmp_obj, tokenized_obj, created_time) + generators.append(self._wait_one_response(tmp_obj)) + rids.append(tmp_obj.rid) + else: + # Batched parallel sampling still follows a conservative path and + # can be slower than duplicating requests explicitly. + if batch_size > 128: + logger.warning( + "Sending a single large batch with parallel sampling (n > 1) has not been well optimized. " + "The performance might be better if you just duplicate the requests n times or use " + "many threads to send them one by one with parallel sampling (n > 1)." + ) + + # Tokenize all requests + objs = [obj[i] for i in range(batch_size)] + tokenized_objs = await self.input_processor.tokenize_batch(objs) + + # Cache the common prefix for parallel sampling + for i in range(batch_size): + tmp_obj = copy.copy(objs[i]) + warmup_obj = prepare_prefix_warmup(tmp_obj, tokenized_objs[i]) + self._send_one_request(tmp_obj, warmup_obj, created_time) + await self._wait_one_response(tmp_obj).__anext__() + + # Expand requests, assign new rids for them, and send them + for i in range(batch_size): + for _ in range(obj.parallel_sample_num): + tmp_obj = copy.copy(objs[i]) + replica_obj = prepare_parallel_sampling_replica( + tmp_obj, tokenized_objs[i] + ) + self._send_one_request(tmp_obj, replica_obj, created_time) + generators.append(self._wait_one_response(tmp_obj)) + rids.append(tmp_obj.rid) + + # Wait for all requests + is_stream = hasattr(obj, "stream") and obj.stream + if not is_stream: + outputs = await asyncio.gather(*(gen.__anext__() for gen in generators)) + yield outputs + else: + rid_to_index = {rid: i for i, rid in enumerate(rids)} + task_map = {asyncio.create_task(gen.__anext__()): gen for gen in generators} + while task_map: + done, _ = await asyncio.wait( + task_map.keys(), return_when=asyncio.FIRST_COMPLETED + ) + + for task in done: + gen = task_map.pop(task) + try: + result = task.result() + result["index"] = rid_to_index[result["meta_info"]["id"]] + yield result + new_task = asyncio.create_task(gen.__anext__()) + task_map[new_task] = gen + except StopAsyncIteration: + pass + + async def flush_cache(self) -> FlushCacheReqOutput: + return (await self.flush_cache_communicator(FlushCacheReqInput()))[0] + + def abort_request(self, rid: str): + if rid not in self.rid_to_state: + return + del self.rid_to_state[rid] + req = AbortReq(rid) + self.engine_core_client.send_to_scheduler.send_pyobj(req) + + async def update_weights_from_disk( + self, + obj: UpdateWeightFromDiskReqInput, + ) -> tuple[bool, str, Any]: + self.auto_create_handle_loop() + + # default the load format to the server_args + if obj.load_format is None: + obj.load_format = self.server_args.load_format + logger.info("Start update_weights. Load format=%s", obj.load_format) + + # Hold the lock if it is not async. This means that weight sync + # cannot run while requests are in progress. + async with self.model_update_lock.writer_lock: + return await self._wait_for_model_update_from_disk(obj) + + async def _wait_for_model_update_from_disk( + self, obj: UpdateWeightFromDiskReqInput + ) -> tuple[bool, str]: + self.engine_core_client.send_to_scheduler.send_pyobj(obj) + self.model_update_result = asyncio.Future() + if not self.server_args.mapping.attn.has_dp: + result = await self.model_update_result + if result.success: + self.served_model_name = obj.model_path + self.server_args.model = obj.model_path + self.server_args.load_format = obj.load_format + self.model_path = obj.model_path + return result.success, result.message, result.num_paused_requests + else: # self.server_args.mapping.has_attn_dp + self.model_update_tmp = [] + result = await self.model_update_result + + all_success = all([r.success for r in result]) + if all_success is True: + self.server_args.model = obj.model_path + self.server_args.load_format = obj.load_format + self.model_path = obj.model_path + all_message = [r.message for r in result] + all_message = " | ".join(all_message) + all_paused_requests = [r.num_paused_requests for r in result] + return all_success, all_message, all_paused_requests + + async def open_session(self, obj: OpenSessionReqInput) -> str | None: + self.auto_create_handle_loop() + + if obj.session_id is None: + obj.session_id = uuid.uuid4().hex + elif obj.session_id in self.session_futures: + return None + + self.engine_core_client.send_to_scheduler.send_pyobj(obj) + + self.session_futures[obj.session_id] = asyncio.Future() + session_id = await self.session_futures[obj.session_id] + del self.session_futures[obj.session_id] + return session_id + + async def close_session(self, obj: CloseSessionReqInput) -> None: + await self.engine_core_client.send_to_scheduler.send_pyobj(obj) + + async def watch_load_thread(self): + # Only for dp_controller when dp_size > 1 + if ( + not self.server_args.mapping.attn.has_dp + or self.server_args.load_balance_method == "round_robin" + ): + return + + while True: + await asyncio.sleep(self.server_args.load_watch_interval) + loads = await self.get_load_communicator(GetLoadReqInput()) + load_udpate_req = WatchLoadUpdateReq(loads=loads) + self.engine_core_client.send_to_scheduler.send_pyobj(load_udpate_req) + + def get_log_request_metadata(self): + max_length = None + skip_names = None + out_skip_names = None + if self.log_requests: + if self.log_requests_level == 0: + max_length = 1 << 30 + skip_names = set( + [ + "text", + "input_ids", + "input_embeds", + "image_data", + "audio_data", + "precomputed_multimodal_inputs", + "input_multi_ids", + ] + ) + out_skip_names = set( + [ + "text", + "output_ids", + ] + ) + elif self.log_requests_level == 1: + max_length = 2048 + elif self.log_requests_level == 2: + max_length = 1 << 30 + else: + raise ValueError( + f"Invalid --log-requests-level: {self.log_requests_level=}" + ) + return max_length, skip_names, out_skip_names + + def configure_logging(self, obj: ConfigureLoggingReq): + if obj.log_requests is not None: + self.log_requests = obj.log_requests + if obj.log_requests_level is not None: + self.log_requests_level = obj.log_requests_level + if obj.dump_requests_folder is not None: + self.dump_requests_folder = obj.dump_requests_folder + if obj.dump_requests_threshold is not None: + self.dump_requests_threshold = obj.dump_requests_threshold + logging.info("Config logging: obj=%r", obj) + self.log_request_metadata = self.get_log_request_metadata() + + # ---- Server lifecycle / health ------------------------------- + # Intent-revealing wrappers around the private ``server_status`` + # field. Callers (notably ``http_server.py``) drive transitions + # through these methods so the ``ServerStatus`` enum and the + # attribute name stay implementation-private. + + def is_server_starting(self) -> bool: + return self.server_status == ServerStatus.Starting + + def mark_server_up(self) -> None: + self.server_status = ServerStatus.Up + + def mark_server_unhealthy(self) -> None: + self.server_status = ServerStatus.UnHealthy + + def drop_request_state(self, rid: str) -> None: + """Discard the per-request state for ``rid`` if present. + + Used by health probes that synthesize a request, await one + token through ``generate_request``, and then need to clean + up the state slot regardless of whether the probe succeeded + or timed out. + """ + self.rid_to_state.pop(rid, None) + + def auto_create_handle_loop(self): + if self.no_create_loop: + return + + self.no_create_loop = True + loop = asyncio.get_event_loop() + self.asyncio_tasks.add( + loop.create_task(print_exception_wrapper(self.handle_loop)) + ) + + # We cannot add signal handler when the tokenizer manager is not in + # the main thread due to the CPython limitation. + if threading.current_thread() is threading.main_thread(): + signal_handler = SignalHandler(self) + loop.add_signal_handler(signal.SIGTERM, signal_handler.signal_handler) + else: + logger.warning( + "Signal handler is not added because the tokenizer manager is " + "not in the main thread. This disables graceful shutdown of the " + "tokenizer manager when SIGTERM is received." + ) + self.asyncio_tasks.add( + loop.create_task(print_exception_wrapper(self.sigterm_watchdog)) + ) + self.asyncio_tasks.add( + loop.create_task(print_exception_wrapper(self.watch_load_thread)) + ) + + async def sigterm_watchdog(self): + while not self.gracefully_exit: + await asyncio.sleep(5) + + # Drain requests + while True: + remain_num_req = len(self.rid_to_state) + logger.info( + "Gracefully exiting... remaining number of requests %s", remain_num_req + ) + if remain_num_req > 0: + await asyncio.sleep(5) + else: + break + + kill_process_tree(os.getpid(), include_parent=True) + sys.exit(0) + + async def handle_loop(self): + """The event loop that handles requests""" + + while True: + recv_obj = await self.engine_core_client.recv_from_detokenizer.recv_pyobj() + self._result_dispatcher(recv_obj) + self.last_receive_tstamp = time.time() + + def _handle_open_session_req_output(self, recv_obj): + self.session_futures[recv_obj.session_id].set_result( + recv_obj.session_id if recv_obj.success else None + ) + + def _handle_update_weights_from_disk_req_output(self, recv_obj): + if not self.server_args.mapping.attn.has_dp: + self.model_update_result.set_result(recv_obj) + else: # self.server_args.mapping.has_attn_dp + self.model_update_tmp.append(recv_obj) + # set future if the all results are received + if len(self.model_update_tmp) == self.server_args.mapping.attn.dp_size: + self.model_update_result.set_result(self.model_update_tmp) + + +async def print_exception_wrapper(func): + """ + Sometimes an asyncio function does not print exception. + We do another wrapper to handle the exception. + """ + try: + await func() + except Exception: + traceback = get_exception_traceback() + logger.error("AsyncLLM hit an exception: %s", traceback) + kill_process_tree(os.getpid(), include_parent=True) + sys.exit(1) + + +class SignalHandler: + def __init__(self, tokenizer_manager): + self.tokenizer_manager = tokenizer_manager + + def signal_handler(self, signum=None, frame=None): + logger.warning( + "SIGTERM received. signum=%r frame=%r. Draining requests and shutting down...", + signum, + frame, + ) + self.tokenizer_manager.gracefully_exit = True + + +T = TypeVar("T") + + +class _Communicator(Generic[T]): + """Note: The communicator now only run up to 1 in-flight request at any time.""" + + def __init__(self, sender, fan_out: int): + self._sender = sender + self._fan_out = fan_out + self._result_event: asyncio.Event | None = None + self._result_values: list[T] | None = None + self._ready_queue: deque[asyncio.Future] = deque() + + async def __call__(self, obj): + ready_event = asyncio.Event() + if self._result_event is not None or len(self._ready_queue) > 0: + self._ready_queue.append(ready_event) + await ready_event.wait() + if self._result_event is not None or self._result_values is not None: + raise RuntimeError("Communicator result state was not reset.") + + if obj: + self._sender.send_pyobj(obj) + + self._result_event = asyncio.Event() + self._result_values = [] + await self._result_event.wait() + result_values = self._result_values + self._result_event = self._result_values = None + + if len(self._ready_queue) > 0: + self._ready_queue.popleft().set() + + return result_values + + def handle_recv(self, recv_obj: T): + self._result_values.append(recv_obj) + if len(self._result_values) == self._fan_out: + self._result_event.set() diff --git a/python/tokenspeed/runtime/engine/collector.py b/python/tokenspeed/runtime/engine/collector.py new file mode 100644 index 0000000..19a4056 --- /dev/null +++ b/python/tokenspeed/runtime/engine/collector.py @@ -0,0 +1,200 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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-processing helpers for the async frontend.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +# Streaming merge policy for the logprob meta_info fields. These are the UNION +# of both dialects' per-position list keys; only one dialect is present per +# request (the renderer emits one), so the union is harmless. +# +# Each value is a per-position list that GROWS as frames arrive, so it must be +# appended rather than overwritten -- hence merged by ``_extend_sequence``. That +# helper is dict-safe: its prefix check (``_is_prefix``) compares elements only +# with ``==``, which both ``dict`` and the ``Logprob`` dataclass support, so the +# entries never need to be hashed or ordered. +# +# ``cumulative_logprob`` is a scalar handled separately (see _SUM_META_KEYS): +# under streaming each frame recomputes it from a fresh dict, so each frame's +# value is the sum of only that frame's positions (a delta). Since the +# per-position ``logprobs`` are appended across frames, the scalar must be +# *summed* (not overwritten) to stay consistent with the appended list. +_APPEND_META_KEYS = { + # vLLM dialect + "logprobs", + # SGLang dialect + "input_token_logprobs", + "output_token_logprobs", + "input_top_logprobs", + "output_top_logprobs", + "input_token_ids_logprobs", + "output_token_ids_logprobs", +} + +# Scalar logprob metadata accumulated by addition across coalesced frames. +_SUM_META_KEYS = { + "cumulative_logprob", +} + + +class RequestOutputCollector: + """Coalesce pending per-request outputs into a single visible response. + + Streaming merges mutate an owned pending dict in place so that N sequential + ``put`` calls cost O(total_delta) instead of O(N * total_delta). The first + merge after a take/reset clones the held output once to detach it from the + producer's reference; subsequent merges extend the cloned lists directly. + """ + + def __init__(self) -> None: + self._pending: dict[str, Any] | None = None + self._pending_owned: bool = False + + def has_pending(self) -> bool: + return self._pending is not None + + def take(self) -> dict[str, Any] | None: + output = self._pending + self._pending = None + self._pending_owned = False + return output + + def put(self, output: dict[str, Any], *, stream: bool) -> None: + if self._pending is None or not stream: + self._pending = output + self._pending_owned = False + return + if not self._pending_owned: + self._pending = self._clone_for_merge(self._pending) + self._pending_owned = True + self._merge_into_pending(output) + + def _merge_into_pending(self, output: dict[str, Any]) -> None: + pending = self._pending + if pending is None: + raise RuntimeError("Cannot merge output without a pending value.") + pending_kind = self._output_kind(pending) + output_kind = self._output_kind(output) + if pending_kind != output_kind: + raise ValueError( + f"Cannot merge different output kinds: {pending_kind} vs {output_kind}" + ) + + if output_kind == "embedding": + # Embedding outputs are latest-wins; drop the owned pending. + self._pending = output + self._pending_owned = False + return + + pending_meta = pending.setdefault("meta_info", {}) + self._merge_meta_info_into(pending_meta, output.get("meta_info") or {}) + + if output_kind == "text" and "text" in output: + pending["text"] = output["text"] + + self._extend_sequence(pending, "output_ids", output.get("output_ids")) + + if "output_multi_ids" in pending or "output_multi_ids" in output: + self._extend_sequence( + pending, "output_multi_ids", output.get("output_multi_ids") + ) + + if "output_extra_info" in output: + pending["output_extra_info"] = output["output_extra_info"] + + def _merge_meta_info_into( + self, pending: dict[str, Any], output: dict[str, Any] + ) -> None: + for key, value in output.items(): + if key == "id": + existing = pending.get("id") + if existing is not None and existing != value: + raise ValueError( + f"Cannot merge outputs for different request ids: " + f"{existing} vs {value}" + ) + pending["id"] = value + continue + if key in _APPEND_META_KEYS: + self._extend_sequence(pending, key, value) + continue + if key in _SUM_META_KEYS: + if value is not None: + pending[key] = (pending.get(key) or 0.0) + value + continue + pending[key] = value + + def _extend_sequence(self, container: dict[str, Any], key: str, value: Any) -> None: + if value is None: + return + existing = container.get(key) + if existing is None: + # Adopt a fresh owned copy so later extends stay private to us. + container[key] = list(value) + return + if not value: + # Follow-up empty list: preserve already-populated values + # (input-logprob producers emit once, then send empty + # lists on subsequent frames). + return + if not existing: + container[key] = list(value) + return + if self._is_prefix(existing, value): + # Cumulative producer: extend with just the tail of `value`. + existing.extend(value[len(existing) :]) + else: + existing.extend(value) + + def _clone_for_merge(self, pending: dict[str, Any]) -> dict[str, Any]: + cloned: dict[str, Any] = dict(pending) + meta = pending.get("meta_info") + if isinstance(meta, dict): + cloned_meta = dict(meta) + for key in _APPEND_META_KEYS: + seq = cloned_meta.get(key) + if isinstance(seq, list): + cloned_meta[key] = list(seq) + cloned["meta_info"] = cloned_meta + for key in ("output_ids", "output_multi_ids"): + seq = cloned.get(key) + if isinstance(seq, list): + cloned[key] = list(seq) + return cloned + + def _output_kind(self, output: dict[str, Any]) -> str: + if "embedding" in output: + return "embedding" + if "text" in output: + return "text" + return "tokens" + + def _is_prefix(self, pending: Sequence[Any], output: Sequence[Any]) -> bool: + pending_len = len(pending) + if pending_len > len(output): + return False + for index in range(pending_len): + if output[index] != pending[index]: + return False + return True diff --git a/python/tokenspeed/runtime/engine/core_client.py b/python/tokenspeed/runtime/engine/core_client.py new file mode 100644 index 0000000..e5ddef6 --- /dev/null +++ b/python/tokenspeed/runtime/engine/core_client.py @@ -0,0 +1,64 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Frontend-side scheduler IPC client for ``AsyncLLM``. + +``EngineCoreClient`` owns the ZMQ context and the two sockets that +``AsyncLLM`` uses to talk to the scheduler subprocess: + +* ``send_to_scheduler`` — ``PUSH`` socket on + ``PortArgs.scheduler_input_ipc_name``; carries tokenized requests, + weight-sync / session / memory-occupation control messages, and the + load-update watcher. +* ``recv_from_detokenizer`` — ``PULL`` socket on + ``PortArgs.tokenizer_ipc_name``; receives ``BatchStrOut`` / + ``BatchTokenIDOut`` / ``BatchEmbeddingOut`` and control-plane replies from + the scheduler. + +Concrete (not ABC): tokenspeed has a single transport (ZMQ in-proc +over ``PortArgs``-provided names) and a single caller (``AsyncLLM``), +so the client stays a plain class. If a second transport ever lands, +promoting this to ``EngineCoreClient(ABC)`` is a purely additive +change. +""" + +import zmq +import zmq.asyncio + +from tokenspeed.runtime.utils import get_zmq_socket +from tokenspeed.runtime.utils.server_args import PortArgs + + +class EngineCoreClient: + """Owns the scheduler-facing ZMQ sockets for ``AsyncLLM``. + + Instantiated once per ``AsyncLLM`` in the front-end process. Socket + attributes are exposed directly so call sites can keep the existing + ``send_pyobj`` / ``recv_pyobj`` ergonomics without a wrapper layer. + """ + + def __init__(self, port_args: PortArgs): + self.context = zmq.asyncio.Context(2) + self.recv_from_detokenizer = get_zmq_socket( + self.context, zmq.PULL, port_args.tokenizer_ipc_name, True + ) + self.send_to_scheduler = get_zmq_socket( + self.context, zmq.PUSH, port_args.scheduler_input_ipc_name, True + ) diff --git a/python/tokenspeed/runtime/engine/data_parallel_controller.py b/python/tokenspeed/runtime/engine/data_parallel_controller.py new file mode 100755 index 0000000..751623d --- /dev/null +++ b/python/tokenspeed/runtime/engine/data_parallel_controller.py @@ -0,0 +1,381 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""A controller that dispatches requests to multiple data parallel workers.""" + +import copy +import multiprocessing as mp +import os +import signal +import threading +from collections import deque +from enum import Enum, auto + +import psutil +import setproctitle +import zmq + +from tokenspeed.runtime.engine.event_loop import run_event_loop +from tokenspeed.runtime.engine.io_struct import ( + BlockReqInput, + TokenizedEmbeddingReqInput, + TokenizedGenerateReqInput, + WatchLoadUpdateReq, +) +from tokenspeed.runtime.engine.request import Req +from tokenspeed.runtime.utils import ( + configure_logger, + get_colorful_logger, + get_zmq_socket, +) +from tokenspeed.runtime.utils.dispatch import TypeBasedDispatcher +from tokenspeed.runtime.utils.exceptions import get_exception_traceback +from tokenspeed.runtime.utils.process import register_usr_signal +from tokenspeed.runtime.utils.server_args import PortArgs, ServerArgs + +logger = get_colorful_logger(__name__) + + +class LoadBalanceMethod(Enum): + """Load balance method.""" + + ROUND_ROBIN = auto() + SHORTEST_QUEUE = auto() + MINIMUM_CACHE_USAGE = 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 DPBudget: + def __init__(self, method: LoadBalanceMethod): + # Use different metrics for load balancing: + # - SHORTEST_QUEUE: by num_reqs (running + waiting) + # - MINIMUM_CACHE_USAGE: by num_pages (page usage) + self.method = method + self.budget_queue = deque() + + def update_budget(self, load_update: WatchLoadUpdateReq): + """Update the budget queue. + + For SHORTEST_QUEUE, use num_reqs instead of num_waiting_reqs to balance decode running batch. + For MINIMUM_CACHE_USAGE, use num_pages as cache usage metric. + """ + # method update_budget and method dispatch happen in the same thread, so clearing budget_queue is safe + self.budget_queue.clear() + + loads = load_update.loads + if not loads: + return + + if self.method == LoadBalanceMethod.MINIMUM_CACHE_USAGE: + metrics = [load.num_pages for load in loads] + else: + metrics = [load.num_reqs for load in loads] + + max_metric = max(metrics) + if all(x == max_metric for x in metrics): + return + + while any(x != metrics[0] for x in metrics): + min_load = min(metrics) + min_indices = [i for i, x in enumerate(metrics) if x == min_load] + second_min_load = min(x for x in metrics if x > min_load) + self.budget_queue.extend( + [loads[i].dp_rank for i in min_indices] * (second_min_load - min_load) + ) + for idx in min_indices: + metrics[idx] = second_min_load + + def dispatch(self): + if self.budget_queue: + return self.budget_queue.popleft() + return None + + +class DataParallelController: + """A controller that dispatches requests to multiple data parallel workers.""" + + def __init__(self, server_args: ServerArgs, port_args: PortArgs) -> None: + # Parse args + self.max_total_num_tokens = None + self.max_req_input_len = None + self.max_num_seqs = None + self.chunked_prefill_size = None + self.max_model_len = 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.mapping.attn.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 + ) + # dp_worker for fixed data dispatch can be set by SINGLE_WORKER_ID environment variable + robin_scheduler = ( + self.round_robin_scheduler + if os.environ.get("SINGLE_WORKER_ID", "-1") == "-1" + else self.single_robin_scheduler + ) + # Dispatch method + self.round_robin_counter = 0 + dispatch_lookup = { + LoadBalanceMethod.ROUND_ROBIN: robin_scheduler, + LoadBalanceMethod.SHORTEST_QUEUE: self.budget_scheduler, + LoadBalanceMethod.MINIMUM_CACHE_USAGE: self.budget_scheduler, + } + self.dispatching = dispatch_lookup[self.load_balance_method] + + # Load balance budget + self.dp_budget = DPBudget(self.load_balance_method) + + # Launch data parallel workers + self.scheduler_procs = [] + self.workers = [None] * server_args.mapping.attn.dp_size + + self.launch_dp_schedulers(server_args, port_args) + + # Workers are already created in launch_dp_schedulers before starting scheduler threads + + if server_args.mapping.has_attn_dp: + self.control_message_step = server_args.mapping.attn.tp_size + else: + self.control_message_step = 1 + + self.init_dispatcher() + + def send_to_all_workers(self, obj): + for worker in self.workers: + worker.send_pyobj(obj) + + def send_control_message(self, obj): + # Send control messages to first worker of tp group + for worker in self.workers[:: self.control_message_step]: + worker.send_pyobj(obj) + + def handle_load_update_req(self, obj): + self.dp_budget.update_budget(obj) + + def init_dispatcher(self): + self._request_dispatcher = TypeBasedDispatcher( + [ + (TokenizedGenerateReqInput, self.dispatching), + (TokenizedEmbeddingReqInput, self.dispatching), + (BlockReqInput, self.send_to_all_workers), + (WatchLoadUpdateReq, self.handle_load_update_req), + ] + ) + self._request_dispatcher.add_fallback_fn(self.send_control_message) + + def launch_dp_schedulers(self, server_args, port_args): + threads = [] + dp_port_args = [] + + # Parse dist_init_addr from port_args to create per-dp-rank ports + # Extract base info from the passed port_args + base_scheduler_port = int(port_args.scheduler_input_ipc_name.split(":")[-1]) + dist_init_host = port_args.scheduler_input_ipc_name.split("//")[1].split(":")[0] + + # port_args.scheduler_input_ipc_name (base_scheduler_port) is used by: + # TokenizerManager -> DataParallelController + # + # For DataParallelController -> Scheduler[dp_rank], we need different ports. + # Following the same logic as PortArgs.init_new with dp_rank parameter: + # scheduler_input_port = port_base + 4 + dp_rank + # Since base_scheduler_port = port_base + 4, we have: + # scheduler_input_port = base_scheduler_port + dp_rank + # + # But we need to avoid conflict with TokenizerManager's port (base_scheduler_port). + # So we start from base_scheduler_port + 1 for dp_rank=0. + + for dp_rank in range(server_args.mapping.attn.dp_size): + # Create port_args for each dp_rank by adjusting scheduler_input_port + # This avoids calling PortArgs.init_new which might use default port + # Use base_scheduler_port + 1 + dp_rank to avoid conflict with TokenizerManager + scheduler_input_port = base_scheduler_port + 1 + dp_rank + tmp_port_args = PortArgs( + tokenizer_ipc_name=port_args.tokenizer_ipc_name, + scheduler_input_ipc_name=f"tcp://{dist_init_host}:{scheduler_input_port}", + nccl_port=port_args.nccl_port, + rpc_ipc_name=port_args.rpc_ipc_name, + metrics_ipc_name=port_args.metrics_ipc_name, + tokenizer_worker_ipc_name=port_args.tokenizer_worker_ipc_name, + ) + dp_port_args.append(tmp_port_args) + + # Bind to scheduler_input_ipc_name BEFORE starting scheduler threads + # This ensures the port is available when scheduler tries to connect + if server_args.node_rank == 0: + self.workers[dp_rank] = get_zmq_socket( + self.context, + zmq.PUSH, + tmp_port_args.scheduler_input_ipc_name, + True, # bind + ) + + if not server_args.mapping.attn.has_dp: + dp_rank_range = range(0, 1) + else: + dp_ranks_per_node = ( + server_args.mapping.attn.dp_size // server_args.mapping.nnodes + ) + dp_rank_range = range( + dp_ranks_per_node * server_args.node_rank, + dp_ranks_per_node * (server_args.node_rank + 1), + ) + for dp_rank in dp_rank_range: + # Create a thread for each worker + thread = threading.Thread( + target=self.launch_tensor_parallel_group, + args=(server_args, dp_port_args[dp_rank], dp_rank), + ) + threads.append(thread) + + # Start all threads + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + return dp_port_args + + def launch_tensor_parallel_group( + self, + server_args: ServerArgs, + port_args: PortArgs, + dp_rank: int, + ): + scheduler_pipe_readers = [] + mapping_template = server_args.mapping + attn_tp_size = mapping_template.attn.tp_size + + if attn_tp_size > mapping_template.nprocs_per_node: + attn_tp_ranks_per_node = attn_tp_size // mapping_template.nnodes + attn_tp_rank_range = range( + attn_tp_ranks_per_node * server_args.node_rank, + attn_tp_ranks_per_node * (server_args.node_rank + 1), + ) + else: + attn_tp_rank_range = range(0, attn_tp_size) + for attn_tp_rank in attn_tp_rank_range: + reader, writer = mp.Pipe(duplex=False) + global_rank = dp_rank * attn_tp_size + attn_tp_rank + + # Create per-rank server_args with rank-initialized mapping + rank_server_args = copy.copy(server_args) + rank_server_args.mapping = copy.deepcopy(mapping_template) + rank_server_args.mapping.rank = global_rank + + proc = mp.Process( + target=run_event_loop, + args=( + rank_server_args, + port_args, + writer, + ), + ) + proc.start() + self.scheduler_procs.append(proc) + scheduler_pipe_readers.append(reader) + # Wait for model to finish loading + scheduler_info = [reader.recv() for reader in scheduler_pipe_readers] + + self.max_total_num_tokens = scheduler_info[0]["max_total_num_tokens"] + self.max_req_input_len = scheduler_info[0]["max_req_input_len"] + self.max_num_seqs = scheduler_info[0]["max_num_seqs"] + self.chunked_prefill_size = scheduler_info[0]["chunked_prefill_size"] + self.max_model_len = scheduler_info[0]["max_model_len"] + + def round_robin_scheduler(self, req: Req): + 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 single_robin_scheduler(self, req): + worker_id = int(os.environ.get("SINGLE_WORKER_ID", "-1")) + if not 0 <= worker_id < self.server_args.mapping.attn.dp_size - 1: + raise ValueError(f"Invalid SINGLE_WORKER_ID:{worker_id}") + self.workers[worker_id].send_pyobj(req) + + def budget_scheduler(self, req): + target_worker = self.dp_budget.dispatch() + if target_worker is None: + self.round_robin_scheduler(req) + else: + 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 + self._request_dispatcher(recv_req) + + +def run_data_parallel_controller_process( + server_args: ServerArgs, + port_args: PortArgs, + pipe_writer, +): + setproctitle.setproctitle("tokenspeed::data_parallel_controller") + configure_logger(server_args) + parent_process = psutil.Process().parent() + register_usr_signal() + + try: + controller = DataParallelController(server_args, port_args) + pipe_writer.send( + { + "status": "ready", + "max_total_num_tokens": controller.max_total_num_tokens, + "max_req_input_len": controller.max_req_input_len, + "max_num_seqs": controller.max_num_seqs, + "chunked_prefill_size": controller.chunked_prefill_size, + "max_model_len": controller.max_model_len, + } + ) + if server_args.node_rank == 0: + controller.event_loop() + for proc in controller.scheduler_procs: + proc.join() + logger.error( + "Scheduler or DataParallelController %s terminated with %s", + proc.pid, + proc.exitcode, + ) + except Exception: + traceback = get_exception_traceback() + logger.error("DataParallelController hit an exception: %s", traceback) + parent_process.send_signal(signal.SIGUSR1) diff --git a/python/tokenspeed/runtime/engine/detokenizer.py b/python/tokenspeed/runtime/engine/detokenizer.py new file mode 100644 index 0000000..84417e1 --- /dev/null +++ b/python/tokenspeed/runtime/engine/detokenizer.py @@ -0,0 +1,346 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Incremental detokenization state machine and helpers. + +This module hosts the pure state machine used by AsyncLLM's inline +detokenizer path. Everything here is tokenizer-agnostic — callers +pass a HuggingFace-shaped tokenizer (with a ``batch_decode`` method) +plus a ``BatchTokenIDOut`` and a mutable ``decode_status`` dict. The +state machine mutates ``decode_status`` in place and returns the +per-request incremental output strings to emit. + +The per-request ``IncrementalDetokenizer`` class wraps a single +``DecodeStatus`` and is the preferred entry point for AsyncLLM; the +batch function ``incremental_decode_batch`` remains as the test +harness driver (``test/runtime/test_detokenizer_parity.py``). +""" + +from __future__ import annotations + +import dataclasses +from collections import OrderedDict, defaultdict +from typing import Any + +from tokenspeed.runtime.engine.io_struct import BatchTokenIDOut +from tokenspeed.runtime.utils.env import envs +from tokenspeed.runtime.utils.text import find_printable_text + +# Maximum number of request states that the detokenizer can hold. +# When exceeded, the oldest entries are evicted. Default: 65536 (1<<16). +DETOKENIZER_MAX_STATES = envs.TOKENSPEED_DETOKENIZER_MAX_STATES.get() + + +@dataclasses.dataclass +class DecodeStatus: + """Per-request incremental decoding state.""" + + decoded_text: str + decode_ids: list[int] + surr_offset: int + read_offset: int + # Offset into ``decoded_text`` that has already been streamed to + # the consumer; the next call emits ``output_str[sent_offset:]``. + sent_offset: int = 0 + + +class LimitedCapacityDict(OrderedDict): + """FIFO-evicting ordered dict used as the detokenizer's request table. + + Only inserting a *new* key at capacity triggers eviction — updating an + existing key is a size-preserving operation and must never drop the + oldest entry. Production detokenizer code writes `self.decode_status[rid] + = s` only on the new-request path, so this guard is defensive for any + future caller that uses the dict for updates. + """ + + def __init__(self, capacity: int, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self.capacity = capacity + + def __setitem__(self, key: Any, value: Any) -> None: + if key not in self and len(self) >= self.capacity: + # Remove the oldest element (first item in the dict) + self.popitem(last=False) + super().__setitem__(key, value) + + +def trim_matched_stop( + output: str | list[int], + finished_reason: dict[str, Any], + no_stop_trim: bool, +) -> str | list[int]: + """Trim a matched stop string or drop a matched stop token. + + If ``no_stop_trim`` is set or ``finished_reason`` is falsy, the + output is returned unchanged. Otherwise: + - When ``matched`` is a ``str`` and ``output`` is also a ``str``, + the output is truncated at the first occurrence of the stop + string. + - When ``matched`` is an ``int`` and ``output`` is a ``list`` + (the raw-token id path), the last id is dropped. + Any other shape combination returns ``output`` unchanged. + """ + if no_stop_trim or not finished_reason: + return output + + matched = finished_reason.get("matched", None) + if not matched: + return output + + # 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): + if not output: + return output + return output[:-1] + return output + + +def decode_grouped_batch( + tokenizer: Any, ids: list[list[int]], recv_obj: BatchTokenIDOut +) -> list[str]: + """Batch-decode requests that disagree on skip/spaces settings. + + Groups requests by ``(skip_special_tokens, spaces_between_special_tokens)`` + so each group can go through a single ``tokenizer.batch_decode`` + call with the correct kwargs, then scatters the results back into + their original positions. + """ + groups: dict[Any, list[Any]] = defaultdict(list) + for i, id in enumerate(ids): + key = ( + recv_obj.skip_special_tokens[i], + recv_obj.spaces_between_special_tokens[i], + ) + groups[key].append((i, id)) + + texts: list[Any] = [None] * len(ids) + + for (skip, spaces), items in groups.items(): + indices, group_ids = zip(*items) + decoded_batch = tokenizer.batch_decode( + group_ids, + skip_special_tokens=skip, + spaces_between_special_tokens=spaces, + ) + for idx, text in zip(indices, decoded_batch): + texts[idx] = text + + return texts + + +def incremental_decode_batch( + tokenizer: Any, + decode_status: dict[str, DecodeStatus], + recv_obj: BatchTokenIDOut, +) -> list[str]: + """Run the incremental detokenizer state machine on a single batch. + + Mutates ``decode_status`` in place: each request's DecodeStatus is + either freshly created or has its decode_ids extended, offsets + advanced, and decoded_text committed. Returns the list of + incremental output strings to emit (one per request in the batch). + + Raises RuntimeError if a request disappears from ``decode_status`` + mid-call, which happens when the capacity-limited dict evicts an + earlier rid during a later rid's assignment in the first loop. + """ + bs = len(recv_obj.rids) + + # Initialize decode status for each request and prepare the + # surr_ids / read_ids slices the tokenizer will decode. + read_ids, surr_ids = [], [] + for i in range(bs): + rid = recv_obj.rids[i] + if rid not in 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], + ) + decode_status[rid] = s + else: + s = decode_status[rid] + s.decode_ids.extend(recv_obj.decode_ids[i]) + + read_ids.append( + 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]) + + all_same = (len(set(recv_obj.skip_special_tokens)) <= 1) and ( + len(set(recv_obj.spaces_between_special_tokens)) <= 1 + ) + if all_same: + surr_texts = 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 = 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], + ) + else: + surr_texts = decode_grouped_batch(tokenizer, surr_ids, recv_obj) + read_texts = decode_grouped_batch(tokenizer, read_ids, recv_obj) + + # Incremental decoding + output_strs: list[str] = [] + for i in range(bs): + try: + s = 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 TOKENSPEED_DETOKENIZER_MAX_STATES environment variable to a bigger value than the default value. " + f"The current value is {DETOKENIZER_MAX_STATES}." + ) + 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 = 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 output_strs + + +class IncrementalDetokenizer: + """Per-request incremental detokenizer wrapping a single ``DecodeStatus``. + + Each instance owns a per-request slice of the state machine that + ``incremental_decode_batch`` runs across an entire batch. The + semantics are byte-for-byte identical to the per-i inner loop of + the batch function for a single-request batch — ``process`` is just + a stateful facade for call sites where one-request-at-a-time + processing is more natural than a shared ``decode_status`` dict. + + Stop authority stays with the scheduler. The ``process`` method + does not return a matched stop string or invent finish reasons — + it only consumes ``finished_reason`` as an input flag exactly + like the batch function does. + """ + + def __init__(self, decoded_text: str = "", read_offset: int = 0) -> None: + self._status = DecodeStatus( + decoded_text=decoded_text, + decode_ids=[], + surr_offset=0, + read_offset=read_offset, + ) + + @property + def status(self) -> DecodeStatus: + """Expose the underlying DecodeStatus for cross-checks / telemetry. + + The returned object is the live mutable state, not a copy. + Callers must not mutate it directly — use ``process`` to advance + the state machine. + """ + return self._status + + def process( + self, + tokenizer: Any, + *, + new_decode_ids: list[int], + finished_reason: dict[str, Any] | None = None, + no_stop_trim: bool = False, + skip_special_tokens: bool = True, + spaces_between_special_tokens: bool = True, + ) -> str: + """Process one frame for this request and return the incremental emit. + + Mutates ``self.status`` in place. Semantically equivalent to one + iteration of the per-i loop in ``incremental_decode_batch`` for a + single-request batch: extend decode_ids with the delta, build + surr_ids/read_ids slices, batch_decode both (single-element + batch), run the partial-UTF-8 deferral / commit machinery, then + emit ``output_str[sent_offset:]``. + """ + s = self._status + s.decode_ids.extend(new_decode_ids) + + read_ids = trim_matched_stop( + s.decode_ids[s.surr_offset :], + finished_reason, + no_stop_trim, + ) + surr_ids = s.decode_ids[s.surr_offset : s.read_offset] + + surr_texts = tokenizer.batch_decode( + [surr_ids], + skip_special_tokens=skip_special_tokens, + spaces_between_special_tokens=spaces_between_special_tokens, + ) + read_texts = tokenizer.batch_decode( + [read_ids], + skip_special_tokens=skip_special_tokens, + spaces_between_special_tokens=spaces_between_special_tokens, + ) + + new_text = read_texts[0][len(surr_texts[0]) :] + if finished_reason 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 = trim_matched_stop( + s.decoded_text + new_text, + finished_reason, + no_stop_trim, + ) + # Incrementally send text. + incremental_output = output_str[s.sent_offset :] + s.sent_offset = len(output_str) + return incremental_output diff --git a/python/tokenspeed/runtime/engine/event_loop.py b/python/tokenspeed/runtime/engine/event_loop.py new file mode 100644 index 0000000..0d0c554 --- /dev/null +++ b/python/tokenspeed/runtime/engine/event_loop.py @@ -0,0 +1,1912 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import faulthandler +import signal +import time +from collections import OrderedDict +from dataclasses import dataclass + +import psutil +import setproctitle +import torch +import torch.distributed as dist +import zmq +from tokenspeed_scheduler import PD, Cache, ExecutionEvent, ForwardEvent, Scheduler + +from tokenspeed.runtime.cache.executor.flat_memory_executor import ( + FlatMemoryExecutor, +) +from tokenspeed.runtime.cache.executor.memory_executor import ( + MemoryExecutor, + MemoryExecutorConfig, +) +from tokenspeed.runtime.cache.transfer.types import CacheKind +from tokenspeed.runtime.configs.model_config import ModelConfig +from tokenspeed.runtime.configs.paged_cache_spec import ( + scheduler_ext_flat_kvcache, + validate_flat_scheduler_config, +) +from tokenspeed.runtime.distributed.process_group_manager import ( + process_group_manager as pg_manager, +) +from tokenspeed.runtime.engine.generation_output_processor import OutputProcesser +from tokenspeed.runtime.engine.memory_occupation import MemoryOccupationController +from tokenspeed.runtime.engine.pause import PauseController +from tokenspeed.runtime.engine.request_handler import RequestHandler +from tokenspeed.runtime.engine.scheduler_utils import ( + advance_forward, + cache_event_from_payload, + cache_event_key, + cache_event_to_payload, + cache_sync_debug_enabled, + make_config, + pool_to_paged_cache_groups, + pool_to_prefix_cache_adjunct_spec, + pop_common_cache_event_payloads, + should_use_overlap_schedule, +) +from tokenspeed.runtime.execution.distributed_initializer import ( + DistributedConfig, + DistributedInitializer, +) +from tokenspeed.runtime.execution.factory import ( + ModelExecutorConfig, + create_model_executor, + create_model_runner, +) +from tokenspeed.runtime.execution.forward_batch_info import ForwardMode +from tokenspeed.runtime.execution.types import ModelExecutionResult +from tokenspeed.runtime.grammar.capturable_grammar import GrammarStepInputs +from tokenspeed.runtime.layers.attention.registry import create_attn_components +from tokenspeed.runtime.metrics.collector import EngineMetrics +from tokenspeed.runtime.pd.decode_executor import DisaggDecodeExecutor +from tokenspeed.runtime.pd.factory import ( + create_kv_transfer, + get_kv_args, +) +from tokenspeed.runtime.pd.kv_events import ( + EventPublisherFactory, + KVEventBatch, + NullEventPublisher, + drain_scheduler_kv_events, + scheduler_kv_events_to_wire_events, +) +from tokenspeed.runtime.pd.mooncake.entities import KVManagerArgs +from tokenspeed.runtime.pd.prefill_executor import DisaggPrefillExecutor +from tokenspeed.runtime.sampling.sampling_params import SamplingParams +from tokenspeed.runtime.utils import ( + configure_logger, + get_colorful_logger, + get_zmq_socket, +) +from tokenspeed.runtime.utils.exceptions import get_exception_traceback +from tokenspeed.runtime.utils.nvtx import nvtx_range +from tokenspeed.runtime.utils.process import register_usr_signal +from tokenspeed.runtime.utils.server_args import PortArgs, ServerArgs +from tokenspeed.runtime.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter + +logger = get_colorful_logger(__name__) + + +def calc_l3_query_hashes(scheduler, tokens: list[int]) -> list[str]: + return scheduler.calc_rolling_hash(tokens, apply_match=True) + + +# Sleep between iterations while frozen (PAUSED_ALL) so the keep-mode pause does +# not busy-spin a CPU core waiting for /resume. +_PAUSED_IDLE_SLEEP_S = 0.001 + + +def _forward_op_executes_model_forward(forward_op, *, is_disagg_decode: bool) -> bool: + """Return whether ``forward_op`` will enter the model forward path. + + On decode-side PD, EXTEND ops only start remote KV receive; the model + forward runs after the remote prefill completes and the scheduler advances + the request into decode. Treating those EXTEND ops as model work makes + idle DP ranks enter dummy collectives that the active rank will not match. + """ + if forward_op is None: + return False + if sum(forward_op.input_lengths) <= 0: + return False + if is_disagg_decode and forward_op.num_extends() > 0: + return False + return True + + +class _NullSender: + """No-op ZMQ sender for non-rank-0 workers.""" + + @staticmethod + def send_pyobj(x): + return None + + +@dataclass(frozen=True) +class DpForwardMetadata: + global_num_tokens: list[int] + global_batch_size: list[int] + global_forward_mode: list[int] + all_decode_or_idle: bool + all_extend: bool + need_idle_forward: bool + + +class EventLoop: + def __init__( + self, + server_args: ServerArgs, + port_args: PortArgs, + gpu_id: int, + attn_tp_rank: int, + dp_rank: int, + global_rank: int, + ) -> None: + # Do not pass server_args further down the stack after this point. + + self.server_args = server_args + self.port_args = port_args + self.gpu_id = gpu_id + self.global_rank = global_rank + + self.model_config = self._load_model_config(server_args.model) + if server_args.speculative_draft_model_path is not None: + draft_model_config = self._load_model_config( + server_args.speculative_draft_model_path, + is_draft_worker=True, + ) + else: + draft_model_config = None + + min_per_gpu_mem = self._init_distributed() + + target, draft = create_model_runner( + server_args, self.model_config, draft_model_config, gpu_id, global_rank + ) + self.use_overlap_schedule = should_use_overlap_schedule( + disable_overlap_schedule=server_args.disable_overlap_schedule, + disaggregation_mode=server_args.disaggregation_mode, + ) + self.overlap_schedule_depth = int(self.use_overlap_schedule) + decode_input_tokens = ( + server_args.speculative_num_draft_tokens + if server_args.speculative_algorithm is not None + else 1 + ) + + ( + attn_backend, + token_to_kv_pool, + draft_attn_backend, + draft_token_to_kv_pool, + self.max_total_num_tokens, + mamba_pool_total_chunks, + mamba_pool, + ) = create_attn_components( + server_args, + self.model_config, + gpu_id, + global_rank, + min_per_gpu_mem, + server_args.enable_memory_saver, + draft_model_config, + decode_input_tokens=decode_input_tokens, + overlap_schedule_depth=self.overlap_schedule_depth, + ) + + num_total_pages = self.max_total_num_tokens // server_args.block_size + hf_config = getattr(self.model_config, "hf_config", None) + text_config = getattr(hf_config, "text_config", None) if hf_config else None + has_mamba = getattr(self.model_config, "mambaish_config", None) is not None or ( + text_config is not None and hasattr(text_config, "mamba2_cache_params") + ) + + mapping = server_args.mapping + # The C++ scheduler's req_pool_idx range is rank-local and 1-based: + # real rows are 1..max_batch_size, row 0 is reserved, and CUDA graph + # padding needs one non-real sink row after the scheduler-owned range. + per_rank_max_batch = server_args.max_num_seqs // max(mapping.attn.dp_size, 1) + req_pool_padding_index = per_rank_max_batch + 1 + + model_executor_config = ModelExecutorConfig.from_server_args( + server_args=server_args, + model_config=self.model_config, + max_req_pool_size=req_pool_padding_index, + gpu_id=gpu_id, + global_rank=global_rank, + num_total_pages=num_total_pages, + overlap_schedule_depth=self.overlap_schedule_depth, + ) + self.model_executor = create_model_executor( + server_args=server_args, + config=model_executor_config, + model_runner=target, + draft_model_runner=draft, + attn_backend=attn_backend, + token_to_kv_pool=token_to_kv_pool, + draft_attn_backend=draft_attn_backend, + draft_token_to_kv_pool=draft_token_to_kv_pool, + mamba_pool=mamba_pool, + ) + + # Reserve one token slot because request validation uses a strict + # ``< max_req_len`` check against the model context length. + self.max_req_input_len = self.model_config.context_len - 1 + self.attn_tp_size = server_args.attn_tp_size or mapping.attn.tp_size + self.world_size = server_args.world_size or mapping.world_size + self.attn_tp_rank = attn_tp_rank + self.attn_tp_cpu_group = pg_manager.get_process_group( + "gloo", server_args.mapping.attn.tp_group + ) + self._pending_cache_event_payloads: OrderedDict[tuple[str, int], dict] = ( + OrderedDict() + ) + # All ranks submit identical cache plans (the C++ scheduler is mirrored), + # so a local in-flight counter mirrors across ranks: if it's 0 here, no + # rank has anything pending. Lets us skip the TP collective in + # _commit_cache_results entirely when nothing is in flight. + self._num_inflight_cache_ops = 0 + self.dp_rank = dp_rank + self.dp_size = mapping.attn.dp_size + self.has_dp = mapping.has_attn_dp + if self.has_dp: + self.world_cpu_group = pg_manager.get_process_group( + "gloo", mapping.world_group + ) + self._dp_local_info = torch.zeros(1, 3, dtype=torch.int32) + self._dp_global_info = torch.zeros(mapping.world_size, 3, dtype=torch.int32) + if not server_args.enable_kvstore: + logger.warning( + "KVStore L2 cache will not be used during normal execution, but it will still be used when retraction happens." + ) + + mamba_l2_host_slots = 0 + if has_mamba and server_args.enable_mamba_l2: + if server_args.mamba_l2_host_slots > 0: + mamba_l2_host_slots = server_args.mamba_l2_host_slots + elif server_args.mamba_l2_host_gb > 0 and mamba_pool is not None: + slot_bytes = int( + mamba_pool.conv_state.shape[0] + * ( + mamba_pool.conv_state[0, 0].nbytes + + mamba_pool.ssm_state[0, 0].nbytes + ) + ) + mamba_l2_host_slots = int( + server_args.mamba_l2_host_gb * (1024**3) // max(slot_bytes, 1) + ) + else: + mamba_l2_host_slots = max( + int(mamba_pool_total_chunks * server_args.mamba_l2_ratio), 1 + ) + + mem_cfg = MemoryExecutorConfig( + layer_num=self.model_config.num_hidden_layers, + page_size=server_args.block_size, + host_ratio=server_args.kvstore_ratio, + host_size_gb=server_args.kvstore_size, + host_parallel_count=max( + int(getattr(server_args.mapping, "nprocs_per_node", 1) or 1), 1 + ), + io_backend=server_args.kvstore_io_backend, + host_layout=server_args.kvstore_mem_layout, + storage_backend=server_args.kvstore_storage_backend, + storage_backend_extra_config=server_args.kvstore_storage_backend_extra_config, + model_name=server_args.model, + enable_mamba_l2=server_args.enable_mamba_l2, + mamba_l2_host_slots=mamba_l2_host_slots, + mamba_l2_layout=server_args.mamba_l2_layout, + mamba_l2_io_backend=server_args.mamba_l2_io_backend, + ) + if scheduler_ext_flat_kvcache() and server_args.enable_kvstore: + if server_args.kvstore_storage_backend is not None: + raise NotImplementedError( + "flat scheduler build (TOKENSPEED_FLAT_KVCACHE) has no L3 " + "storage tier yet; unset --kvstore-storage-backend." + ) + self.memory_executor = FlatMemoryExecutor( + device_pool=token_to_kv_pool, + host_ratio=server_args.kvstore_ratio, + host_size_gb=server_args.kvstore_size, + ) + num_host_pages = self.memory_executor.num_host_pages + elif not token_to_kv_pool.supports_hierarchical_kv_cache: + if server_args.enable_kvstore: + raise NotImplementedError( + "This KV cache pool does not support hierarchical cache " + "(kvstore); pass --disable-kvstore." + ) + self.memory_executor = None + num_host_pages = 0 + else: + self.memory_executor = MemoryExecutor( + device_pool=token_to_kv_pool, + config=mem_cfg, + is_dp_attention_enabled=self.has_dp, + tp_group=self.attn_tp_cpu_group, + draft_device_pool=draft_token_to_kv_pool, + mamba_pool=mamba_pool, + ) + num_host_pages = self.memory_executor.host_pool.page_num + + # Flat host tier acks loadbacks (LoadBackDoneEvent), so they join the + # inflight accounting in _submit_cache_ops; radix loadbacks never ack. + self._loadback_acks_expected = getattr( + self.memory_executor, "emits_loadback_acks", False + ) + + self._kv_events_enabled = ( + EventPublisherFactory.is_enabled(server_args.kv_events_config) + and attn_tp_rank == 0 + ) + + if has_mamba and server_args.max_mamba_cache_size is None: + logger.info( + f"Mamba radix cache enabled without explicit max_mamba_cache_size. " + f"Auto-derived mamba_pool_total_chunks={mamba_pool_total_chunks} " + f"(ratio={server_args.mamba_full_memory_ratio})." + ) + + # Adjunct enabled only when pool opts in AND prefix-caching switch is on. + paged_cache_groups = pool_to_paged_cache_groups(token_to_kv_pool) + validate_flat_scheduler_config( + flat_kvcache_ext=scheduler_ext_flat_kvcache(), + paged_cache_groups=paged_cache_groups, + attn_backend=attn_backend, + kv_pool=token_to_kv_pool, + speculative_algorithm=server_args.speculative_algorithm, + ) + self._paged_cache_groups = paged_cache_groups + prefix_cache_adjunct = None + required_groups = token_to_kv_pool.prefix_cache_required_group_ids + if required_groups is not None and server_args.enable_prefix_caching: + prefix_cache_adjunct = pool_to_prefix_cache_adjunct_spec(required_groups) + scheduler_cfg = make_config( + num_device_pages=self.max_total_num_tokens // server_args.block_size, + max_scheduled_tokens=server_args.chunked_prefill_size, + max_batch_size=per_rank_max_batch, + page_size=server_args.block_size, + num_host_pages=num_host_pages, + disable_l2_cache=not server_args.enable_kvstore, + enable_l3_storage=server_args.kvstore_storage_backend is not None, + prefetch_threshold=4, # Keep this hard-coded until it becomes configurable. + role=server_args.disaggregation_mode, + enable_kv_cache_events=self._kv_events_enabled, + decode_input_tokens=decode_input_tokens, + overlap_schedule_depth=self.overlap_schedule_depth, + disable_prefix_cache=not server_args.enable_prefix_caching, + enable_mamba=has_mamba, + mamba_cache_chunk_size=server_args.mamba_cache_chunk_size, + mamba_pool_total_chunks=mamba_pool_total_chunks, + enable_mamba_l2=server_args.enable_mamba_l2, + mamba_l2_host_slots=mamba_l2_host_slots, + paged_cache_groups=paged_cache_groups, + enable_mixed_prefill_decode=server_args.enable_mixed_batch, + prefix_cache_adjunct=prefix_cache_adjunct, + ) + logger.info( + "Scheduler config: block_size=%s num_device_pages=%s " + "max_scheduled_tokens=%s decode_input_tokens=%s " + "overlap_schedule_depth=%s disable_l2_cache=%s " + "max_batch_size=%s (global max_num_seqs=%s, dp_size=%s) " + "mamba_pool_total_chunks=%s enable_mamba=%s " + "disable_prefix_cache=%s paged_cache_groups=%s", + scheduler_cfg.block_size, + scheduler_cfg.num_device_pages, + scheduler_cfg.max_scheduled_tokens, + scheduler_cfg.decode_input_tokens, + scheduler_cfg.overlap_schedule_depth, + scheduler_cfg.disable_l2_cache, + scheduler_cfg.max_batch_size, + server_args.max_num_seqs, + self.dp_size, + mamba_pool_total_chunks, + has_mamba, + scheduler_cfg.disable_prefix_cache, + [group.group_id for group in paged_cache_groups], + ) + self.scheduler = Scheduler(scheduler_cfg) + token_to_kv_pool.bind_paged_cache_scheduler(self.scheduler) + if attn_tp_rank == 0: + self.kv_event_publisher = EventPublisherFactory.create( + server_args.kv_events_config, + attn_dp_rank=dp_rank, + ) + else: + self.kv_event_publisher = NullEventPublisher(attn_dp_rank=dp_rank) + + self._init_interprocess_comm() + + # Pause/resume control state. Shared with the request handler, which + # drives the control-request side; the event loop reads the gate. + self._pause = PauseController(self.send_to_tokenizer) + + # GPU-memory data plane (release/resume_memory_occupation). Reuses the + # pause controller's drain machinery; frees memory via the memory-saver + # adapter once the scheduler drains. See memory_occupation.py. + # Releasing KV is only safe if any prefix cache it backs can be cleared: + # either prefix caching is off, or the scheduler exposes a reset. Decide + # once here (static config) and let the controller reject unsafe releases. + kv_cache_release_allowed = ( + not self.server_args.enable_prefix_caching + or callable(getattr(self.scheduler, "reset_prefix_cache", None)) + ) + self._memory = MemoryOccupationController( + send_func=self.send_to_tokenizer, + pause_controller=self._pause, + adapter=TorchMemorySaverAdapter.create( + enable=self.server_args.enable_memory_saver + ), + enabled=self.server_args.enable_memory_saver, + reset_caches_fn=self._reset_caches_for_release, + kv_repair_fn=self._kv_repair_after_wake, + kv_cache_release_allowed=kv_cache_release_allowed, + ) + + self.metrics = EngineMetrics( + labels={ + "model_name": server_args.served_model_name, + "app_key": server_args.app_key or "", + "dp_rank": str(dp_rank), + }, + enabled=( + server_args.enable_metrics + and attn_tp_rank == 0 + and "prometheus" in (server_args.metrics_reporters or []) + ), + ) + + self.request_handler = RequestHandler( + server_args=self.server_args, + hf_eos_token_id=self.model_config.hf_eos_token_id, + max_req_len=self.model_config.context_len - 1, + vocab_size=self.model_config.vocab_size, + recv_func=self.recv_from_tokenizer, + send_func=self.send_to_tokenizer, + get_load_fn=self._get_load, + architectures=self.model_config.hf_config.architectures, + pause_controller=self._pause, + memory_controller=self._memory, + ) + + self.output_processor = OutputProcesser( + send_to_tokenizer=self.send_to_tokenizer, + attn_tp_rank=attn_tp_rank, + spec_algorithm=self.server_args.speculative_algorithm, + spec_num_tokens=( + self.server_args.speculative_num_draft_tokens + if self.server_args.speculative_algorithm is not None + else None + ), + stream_interval=self.server_args.stream_interval, + enable_log_request_stats=self.server_args.enable_log_request_stats, + metrics=self.metrics, + ) + self.prefetch_threshold = scheduler_cfg.prefetch_threshold + + if server_args.disaggregation_mode != "null": + kv_args = get_kv_args( + global_rank, + global_rank, + server_args.disaggregation_ib_device, + token_to_kv_pool, + draft_token_to_kv_pool, + mamba_pool, + ) + pd_manager_args = KVManagerArgs( + bootstrap_port=server_args.disaggregation_bootstrap_port, + dist_init_addr=server_args.dist_init_addr, + world_size=server_args.world_size or mapping.world_size, + dp_size=server_args.data_parallel_size or mapping.attn.dp_size, + attn_tp_rank=attn_tp_rank, + attn_dp_rank=dp_rank, + is_mla_backend=False, + draft_is_mla_backend=False, + enable_metrics=False, + enable_mla_l1_5_cache=server_args.enable_mla_l1_5_cache, + served_model_name=server_args.served_model_name, + app_key=server_args.app_key, + metrics_reporters=server_args.metrics_reporters, + enable_dp_attention=self.has_dp, + ) + self.kv_transfer = create_kv_transfer( + mode=server_args.disaggregation_mode, + backend=server_args.disaggregation_transfer_backend, + args=pd_manager_args, + kv_args=kv_args, + gloo_group=self.attn_tp_cpu_group, + page_size=token_to_kv_pool.page_size, + ) + self._setup_pd_layerwise_transfer( + server_args.disaggregation_layerwise_interval + ) + # EPD: a multimodal prefill node is also the encode->prefill embedding + # SINK (independent of kv_transfer, its P->D KV source) -- it receives + # each image's embedding from encode workers over Mooncake so the + # prefill skips the vision tower. The admission controller owns the + # receive jobs, the rank-synced admission drain, and the optional NCCL + # row-shard reassembly; None for decode/encode/text-only nodes. + from tokenspeed.runtime.pd.epd.prefill_receiver import ( + make_epd_prefill_admission, + ) + + self.epd_admission = make_epd_prefill_admission( + server_args, + global_rank, + model_config=self.model_config, + model_executor=self.model_executor, + mapping=mapping, + attn_tp_rank=self.attn_tp_rank, + attn_tp_size=self.attn_tp_size, + attn_tp_cpu_group=self.attn_tp_cpu_group, + pg_manager=pg_manager, + ) + # Staged EPD request payloads (request_id -> (spec, state, bootstrap)), + # held here while the controller (rid-keyed, like kv_transfer) runs the + # async receive; popped in _drain_ready_epd_embeddings on admit/abort. + self._epd_staged: dict = {} + else: + self.kv_transfer = None + self.epd_admission = None + self._epd_staged: dict = {} + + def _setup_pd_layerwise_transfer(self, interval: int) -> None: + if not isinstance(self.kv_transfer, DisaggPrefillExecutor): + return + if interval <= 0: + return + + from tokenspeed.runtime.pd.utils import StepCounter + + step_counter = StepCounter(self.model_executor.device, self.gpu_id) + self.model_executor.attn_backend.register_step_counter(step_counter) + if self.model_executor.draft_attn_backend is not None: + self.model_executor.draft_attn_backend.register_step_counter(step_counter) + self.kv_transfer.register_layerwise_step_counter(step_counter, interval) + + def _is_epd_request(self, state) -> bool: + """True iff this request's images are encode-routed (smg injected per-image + encode handshakes) -- it must wait for its embeddings (staged via the EPD + admission controller, polled in _drain_ready_epd_embeddings) before being + scheduled. Caller guards on self.epd_admission (only a multimodal prefill + node has one); everything else admits immediately. + """ + mm = getattr(state, "multimodal_inputs", None) + return mm is not None and any( + getattr(it, "encode_handshake", None) for it in mm.mm_items + ) + + def _assert_epd_embeddings_received(self, multimodal_context) -> None: + """EPD invariant: every handshaked item is filled with its embedding by the + async EPD admission drain (EpdPrefillAdmission.drain) BEFORE admission, so by + it is already encoded. This is a defensive check, not a receive: a handshaked + item that reached the forward un-received leaked past async admission (the + only EPD admission path) -- fail loud instead of running the tower or + publishing shard-only rows. No-op for non-EPD / text-only requests. + """ + if ( + self.epd_admission is None + or multimodal_context is None + or not multimodal_context.has_extend_inputs() + ): + return + for mm in multimodal_context.mm_inputs: + if mm is None: + continue + missing = [ + i + for i, item in enumerate(mm.mm_items) + if getattr(item, "encode_handshake", None) is not None + and item.encoded is None + ] + if missing: + raise RuntimeError( + f"EPD: handshaked items {missing} reached the prefill forward " + "un-received; they must be admitted via the EPD admission drain" + ) + + def _drain_ready_epd_embeddings(self) -> None: + """Admit EPD requests whose async embedding receives completed this cycle. + + The EpdPrefillAdmission controller DECIDES (poll + rank-lockstep MIN + all-reduce + reassemble) and returns (admitted, failed); here we ACT on + those decisions with the EventLoop's collaborators -- register/abort the + P->D sender, submit admitted requests, finish failed ones. No-op (and no + collective) on non-EPD nodes. + """ + if self.epd_admission is None: + return + # Pause gate: withhold EPD admission while paused, mirroring the non-EPD + # admit_blocked gate -- else the drain below would submit and RUN reassembled + # specs during the pause. Staged receives wait in _pending until resume. + # Rank-safe: admit_blocked is rank-identical, so all ranks skip together. + if self._pause.admit_blocked: + return + admitted_ids, failed_ids = self.epd_admission.drain() + for rid in failed_ids: + spec, state, bootstrap = self._epd_staged.pop(rid) + # Signal the dual-dispatched decode that this request failed so its KV + # receiver fails (FailedEvent -> _process_kv_transfer_events abort) + # instead of waiting forever for KV the prefill will never send. The + # prefill never registered a P->D sender (deferred to admission), so the + # decode has no other reliable way to learn (heartbeat only trips on a + # dead prefill /health). Best-effort: only reaches decodes that already + # pre-allocated. + if ( + isinstance(self.kv_transfer, DisaggPrefillExecutor) + and bootstrap is not None + ): + try: + self.kv_transfer.abort(rid, bootstrap) + except Exception as exc: # never let it wedge the loop + logger.warning( + "EPD abort->decode signal failed for rid=%s: %s", + rid, + exc, + ) + state.set_finish_with_abort("EPD embedding receive failed or timed out") + self.output_processor.publish_finished_at_admission(rid, state) + admitted_specs = [] + for rid in admitted_ids: + spec, state, bootstrap = self._epd_staged.pop(rid) + # Aborted mid-receive (no abort path, so drain still returns it admitted): + # don't register the P->D sender or submit -- that runs a wasted forward + # and leaks the sender. Stream its finish instead. + if state.finished: + self.output_processor.publish_finished_at_admission(rid, state) + continue + # Register the P->D sender now (deferred from admission) -- the request + # is about to enter the scheduler. + if self.kv_transfer is not None: + self.kv_transfer.register(rid, bootstrap) + admitted_specs.append(spec) + if admitted_specs: + self.scheduler.submit_requests(admitted_specs) + elif self.epd_admission.has_pending(): + # Nothing advanced this cycle but requests are still receiving; yield the + # GIL so the Python daemon transfer/recv threads run (rank-consistent: + # admitted/leftover are rank-identical here). + time.sleep(0.0005) + + def _commit_cache_results(self) -> None: + if self.memory_executor is None: + return + cache_results = self.memory_executor.poll_results() + self._num_inflight_cache_ops -= len(cache_results) + for event in cache_results: + payload = cache_event_to_payload(event) + self._pending_cache_event_payloads[cache_event_key(payload)] = payload + + # The gather below is a collective, but cache-op completion is async and + # not lock-step across ranks, so local state (_num_inflight_cache_ops / + # _pending_cache_event_payloads) diverges transiently. A rank-local skip + # would let some ranks gather while others return, deadlocking the group. + # Agree on the skip via a cheap single-int all_reduce. + # NOTE: For non-DFLASH algorithms, cache ops are deterministic across + # ranks, so the local short-circuit is safe and avoids collective overhead. + local_has_work = bool( + self._num_inflight_cache_ops != 0 or self._pending_cache_event_payloads + ) + if self.server_args.speculative_algorithm == "DFLASH": + if not self._cache_group_has_work(local_has_work): + return + else: + if not local_has_work: + return + + ready_payloads = self._pop_ready_cache_event_payloads() + if not ready_payloads: + return + logger.debug( + "[cache_poll] got %s synchronized results, advancing scheduler", + len(ready_payloads), + ) + ec = ExecutionEvent() + for payload in ready_payloads: + e = cache_event_from_payload(payload) + logger.debug( + "[cache_poll] event: op_id=%s success=%s type=%s request_id=%s", + e.op_id, + e.success, + type(e).__name__, + getattr(e, "request_id", "N/A"), + ) + ec.add_event(e) + self.scheduler.advance(ec) + logger.debug("[cache_poll] scheduler.advance() done") + self._publish_scheduler_kv_events() + + def _publish_scheduler_kv_events(self) -> None: + raw_events = drain_scheduler_kv_events( + self.scheduler, + enabled=self._kv_events_enabled, + ) + if not raw_events: + return + + events = scheduler_kv_events_to_wire_events(raw_events) + if not events: + return + + self.kv_event_publisher.publish( + KVEventBatch(ts=time.time(), events=events, attn_dp_rank=self.dp_rank) + ) + + def _cache_group_has_work(self, local_has_work: bool) -> bool: + """Whether ANY attn-tp rank has cache work this step (unanimous via a + single-int MAX all_reduce, far cheaper than the payload gather it + guards). Deciding from rank-local state alone deadlocks the group; see + _commit_cache_results. + + Args: + local_has_work: This rank's view of whether any cache op is in + flight or any polled payload awaits commit. + + Returns: + ``True`` if any rank has work (all must gather); ``False`` only when + every rank is idle. + """ + if self.attn_tp_size == 1: + return local_has_work + flag = torch.tensor([1 if local_has_work else 0], dtype=torch.int32) + dist.all_reduce(flag, op=dist.ReduceOp.MAX, group=self.attn_tp_cpu_group) + return bool(flag.item()) + + def _pop_ready_cache_event_payloads(self) -> list[dict]: + local_payloads = list(self._pending_cache_event_payloads.values()) + if self.attn_tp_size == 1: + ready_payloads = local_payloads + else: + gathered_payloads = [None] * self.attn_tp_size + dist.all_gather_object( + gathered_payloads, + local_payloads, + group=self.attn_tp_cpu_group, + ) + ready_payloads = pop_common_cache_event_payloads(gathered_payloads) + if self.attn_tp_rank == 0 and cache_sync_debug_enabled(): + pending_ops = [ + [(payload["kind"], payload["op_id"]) for payload in rank_payloads] + for rank_payloads in gathered_payloads + ] + if len({tuple(rank_ops) for rank_ops in pending_ops}) > 1: + logger.info( + "[cache_sync] rank=%s pending_ops=%s ready_ops=%s", + self.global_rank, + pending_ops, + [ + (payload["kind"], payload["op_id"]) + for payload in ready_payloads + ], + ) + + for payload in ready_payloads: + self._pending_cache_event_payloads.pop(cache_event_key(payload), None) + return ready_payloads + + def _dispatch_forward( + self, + forward_op, + sampling_params_list, + execution_plan, + dp_metadata=None, + stats=None, + grammar_inputs=None, + ): + """Execute one forward step; return (results, on_first_token). + + results is None when the step produces no model output (Path 2/3). + Both event_loop and event_loop_overlap call this method; they differ + only in *when* they call post_process on the returned results. + + Path 1 — no PD: run forward, return (results, None) + Path 2 — decode, extend: trigger RDMA receive, return (None, None) + Path 3 — prefill, decode: send KV to decode side, return (None, None) + Path 4 — prefill, extend: run prefill forward, return (results, on_first_token) + """ + if stats is None: + stats = {} + dp_global_num_tokens = ( + dp_metadata.global_num_tokens if dp_metadata is not None else None + ) + dp_global_bs = ( + dp_metadata.global_batch_size if dp_metadata is not None else None + ) + dp_all_decode_or_idle = ( + dp_metadata.all_decode_or_idle if dp_metadata is not None else False + ) + dp_all_extend = dp_metadata.all_extend if dp_metadata is not None else False + multimodal_context = self._get_multimodal_context_for_forward(forward_op) + + self.model_executor.update_block_table(forward_op) + + if self.kv_transfer is None: + # Path 1: normal (no disaggregation) + self.model_executor.reset_valid_cache_length(forward_op) + return ( + self.model_executor.execute_forward_op_with_log( + forward_op, + sampling_params_list, + dp_global_num_tokens=dp_global_num_tokens, + dp_global_bs=dp_global_bs, + dp_all_decode_or_idle=dp_all_decode_or_idle, + dp_all_extend=dp_all_extend, + grammar_inputs=grammar_inputs, + multimodal_context=multimodal_context, + **stats, + ), + None, + ) + + elif isinstance(self.kv_transfer, DisaggDecodeExecutor): + # Decode node + if forward_op.num_extends() > 0: + # Path 2: new requests waiting for remote KV — trigger RDMA receive + self.kv_transfer.reset_valid_cache_length( + forward_op, + self.model_executor.runtime_states, + self.model_executor.execution_stream, + self.model_executor.device, + ) + self.kv_transfer.execute(forward_op) + self.model_executor.reset_remote_prefill_mamba_inputs(forward_op) + return None, None + else: + # Path 3b: decode batch — normal forward + self.model_executor.reset_valid_cache_length(forward_op) + return ( + self.model_executor.execute_forward_op_with_log( + forward_op, + sampling_params_list, + dp_global_num_tokens=dp_global_num_tokens, + dp_global_bs=dp_global_bs, + dp_all_decode_or_idle=dp_all_decode_or_idle, + dp_all_extend=dp_all_extend, + multimodal_context=multimodal_context, + **stats, + ), + None, + ) + + else: + # Prefill node (only reached from event_loop, never event_loop_overlap) + if not isinstance(self.kv_transfer, DisaggPrefillExecutor): + raise TypeError("kv_transfer must be a DisaggPrefillExecutor.") + if forward_op.num_extends() == 0: + # Path 3: all prefill done — send KV to decode side + self.kv_transfer.execute(forward_op) + return None, None + else: + # Path 4: extend batch — run prefill forward + self.model_executor.reset_valid_cache_length(forward_op) + self.kv_transfer.prepare_prefill(forward_op) + # EPD invariant: handshaked items are filled by the async + # EPD admission drain before admission; assert none reached + # the forward un-received (no-op for non-EPD / text-only requests). + self._assert_epd_embeddings_received(multimodal_context) + return ( + self.model_executor.execute_forward_op_with_log( + forward_op, + sampling_params_list, + dp_global_num_tokens=dp_global_num_tokens, + dp_global_bs=dp_global_bs, + dp_all_decode_or_idle=dp_all_decode_or_idle, + dp_all_extend=dp_all_extend, + grammar_inputs=grammar_inputs, + multimodal_context=multimodal_context, + capture_next_input_ids=True, + **stats, + ), + self.kv_transfer.store_prefill_token, + ) + + def _get_multimodal_context_for_forward(self, forward_op): + if not self.model_config.is_multimodal_active: + return None + + num_extends = forward_op.num_extends() + mm_inputs = [] + has_mm = False + for index, rid in enumerate(forward_op.request_ids): + state = self.output_processor.rid_to_state.get(rid) + if state is not None and index < num_extends: + state.maybe_extend_multimodal_mrope_positions() + item = getattr(state, "multimodal_inputs", None) if state else None + mm_inputs.append(item) + has_mm = has_mm or item is not None + if not has_mm: + return None + + from tokenspeed.runtime.multimodal.inputs import MultimodalForwardContext + + return MultimodalForwardContext( + mm_inputs=mm_inputs, + extend_prefix_lens=list(forward_op.extend_prefix_lens), + extend_seq_lens=list(forward_op.input_lengths[:num_extends]), + ) + + def _build_mamba_layerwise_cow( + self, execution_plan, forward_op + ) -> dict[int, list[int]]: + if forward_op is None: + return {} + loaded_mamba_slots: set[int] = set() + for cache_op in execution_plan.cache: + if not isinstance(cache_op, Cache.LoadBackOp): + continue + dst_by_kind = getattr(cache_op, "dst_pages_by_kind", None) + if dst_by_kind is None: + dst_groups = getattr(cache_op, "dst_pages", []) + else: + dst_groups = dst_by_kind.get(CacheKind.MAMBA.value, []) + for dst_pages in dst_groups: + loaded_mamba_slots.update(int(page) for page in dst_pages) + if not loaded_mamba_slots: + return {} + + cow_src_indices = getattr(forward_op, "mamba_cow_src_indices", None) + working_indices = getattr(forward_op, "mamba_pool_indices", None) + if cow_src_indices is None or working_indices is None: + return {} + + cow_by_src: dict[int, list[int]] = {} + for cow_src, working in zip(list(cow_src_indices), list(working_indices)): + cow_src = int(cow_src) + working = int(working) + if cow_src < 0 or working < 0 or cow_src not in loaded_mamba_slots: + continue + cow_dsts = cow_by_src.setdefault(cow_src, []) + if working not in cow_dsts: + cow_dsts.append(working) + return cow_by_src + + def _submit_cache_ops(self, execution_plan) -> None: + if self.memory_executor is None: + return + forward_op = self._get_forward_op(execution_plan) + mamba_layerwise_cow = self._build_mamba_layerwise_cow( + execution_plan, forward_op + ) + if mamba_layerwise_cow: + self.model_executor.set_layerwise_mamba_cow_done(mamba_layerwise_cow) + self.memory_executor.set_mamba_layerwise_cow(mamba_layerwise_cow) + self.memory_executor.submit_plan(execution_plan) + for op in execution_plan.cache: + if isinstance(op, Cache.WriteBackOp): + self._num_inflight_cache_ops += len(op.op_ids) + elif isinstance(op, Cache.LoadBackOp): + # Radix loadbacks are fire-and-forget (no ack, nothing in + # flight); the flat host tier acks one LoadBackDone per op_id. + if self._loadback_acks_expected: + self._num_inflight_cache_ops += len(op.op_ids) + elif isinstance(op, (Cache.PrefetchOp, Cache.BackUpOp)): + self._num_inflight_cache_ops += 1 + else: + raise ValueError(f"unsupported cache op kind: {type(op).__name__}") + self._setup_layerwise_loadback(execution_plan) + + def _setup_layerwise_loadback(self, execution_plan) -> None: + host_exec = getattr(self.memory_executor, "host_exec", None) + available_pools = ( + getattr(host_exec, "pools", {}) if host_exec is not None else {} + ) + consumer_indices_by_kind: dict[CacheKind, list[int]] = { + kind: [] for kind in available_pools + } + for cache_op in execution_plan.cache: + if isinstance(cache_op, Cache.LoadBackOp): + for op_id in cache_op.op_ids: + for kind in consumer_indices_by_kind: + producer_idx = self.memory_executor.get_producer_index( + kind, op_id + ) + if ( + producer_idx is not None + and producer_idx not in consumer_indices_by_kind[kind] + ): + consumer_indices_by_kind[kind].append(producer_idx) + for kind, consumer_indices in consumer_indices_by_kind.items(): + self.memory_executor.set_consumer( + kind, consumer_indices if consumer_indices else -1 + ) + + def _flush_mamba_retract_states(self, forward_op) -> None: + """Copy draft->working mamba states when retract occurred (no forward scheduled).""" + if forward_op is not None: + return + if self.model_executor.drafter is None: + return + if self.model_executor.runtime_states.mamba_pool is None: + return + self.model_executor.flush_mamba_draft_to_working_on_retract() + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def _load_model_config( + self, model_path: str, is_draft_worker: bool = False + ) -> ModelConfig: + server_args = self.server_args + quantization = server_args.quantization + if is_draft_worker: + quantization = server_args.speculative_draft_model_quantization + return ModelConfig( + model_path, + trust_remote_code=server_args.trust_remote_code, + revision=server_args.revision, + context_length=server_args.max_model_len, + model_override_args=server_args.hf_overrides, + dtype=server_args.dtype, + quantization=quantization, + server_args=server_args, + is_draft_worker=is_draft_worker, + ) + + def _init_distributed(self) -> float: + max_num_input_tokens = ( + self.server_args.chunked_prefill_size + if self.server_args.chunked_prefill_size > 0 + else self.server_args.max_prefill_tokens + self.server_args.max_model_len + ) + distributed_config = DistributedConfig.from_server_args( + server_args=self.server_args, + port_args=self.port_args, + gpu_id=self.gpu_id, + global_rank=self.global_rank, + hidden_size=self.model_config.hidden_size, + max_num_tokens=max_num_input_tokens, + ) + return DistributedInitializer.initialize(distributed_config) + + def _init_interprocess_comm(self): + context = zmq.Context(2) + if self.attn_tp_rank == 0: + self.recv_from_tokenizer = get_zmq_socket( + context, zmq.PULL, self.port_args.scheduler_input_ipc_name, False + ) + self.send_to_tokenizer = get_zmq_socket( + context, zmq.PUSH, self.port_args.tokenizer_ipc_name, False + ) + else: + self.recv_from_tokenizer = None + self.send_to_tokenizer = _NullSender() + + # ------------------------------------------------------------------ + # Shared step helpers + # ------------------------------------------------------------------ + + def _reap_or_keep_buffered_spec(self, spec) -> bool: + """Resolve a buffered spec on resume; return True if it should be admitted. + + A buffered spec was already registered in ``rid_to_state`` before it was + withheld, so if it was aborted while paused it never reached the + scheduler and the forward path can never reap it. Handle that here: + + - state missing -> already published and reaped; drop silently. + - state finished -> aborted in place. Stream a terminating finish for + pause-initiated aborts (the passive client is still waiting) and drop + the registered state so the rid does not leak; client-initiated aborts + already tore down their own state, so just reap. + - otherwise -> still live; admit it. + """ + state = self.output_processor.rid_to_state.get(spec.request_id) + if state is None: + return False + if state.finished: + self.output_processor.reap_finished_orphan(spec.request_id, state) + return False + return True + + def _process_new_requests(self): + recv_reqs = self.request_handler.recv_reqs() + # Snapshot the pause state before dispatch: process_requests may flip it + # mid-batch. If it was not blocked before but is after, a pause control + # message was processed in this very batch — which is what makes the + # FIFO edge below detectable (see TODO(pause-fifo)). + pause_blocked_before = self._pause.admit_blocked + new_req_specs, new_req_states, bootstrap_infos, abort_rids = ( + self.request_handler.process_requests(recv_reqs) + ) + # Sweep TTL-expired abort markers every iteration. Without this + # the map only gets cleaned inside ``mark_abort``, so a burst of + # stale-cancel traffic followed by silence leaves the last batch + # of entries sitting past their TTL (and potentially re-aborting + # reused rids). Amortized O(1): expired entries are always at + # the front of the insertion-ordered dict. + self.output_processor.sweep_pending_aborts() + # Abort both registered and grammar-queued requests. Without the + # grammar_manager.mark_abort call, a request aborted mid-compile + # would finish compiling and get admitted before being noticed. + grammar_manager = self.request_handler.grammar_manager + for rid in abort_rids: + self.output_processor.mark_abort(rid) + grammar_manager.mark_abort(rid) + + # A pause(mode="abort") cancels every in-flight request through the same + # marker path as a client abort; they finish on their next scheduled + # step, then the drain check resolves the pause reply. + if self._pause.consume_abort_all(): + for rid in list(self.output_processor.rid_to_state.keys()): + # notify_client=True: pause aborts a passive client's request, + # so it must receive a terminating finish (unlike a client abort). + self.output_processor.mark_abort(rid, notify_client=True) + grammar_manager.mark_abort(rid) + + # abort/wait also cancel requests still compiling in the grammar queue: + # they are not yet in rid_to_state or the scheduler, so the sweep above + # and the drain check both miss them. A finished state makes the next + # get_ready_grammar_requests pass publish them instead of admitting, so + # they never run under post-resume weights or strand the drain. + if self._pause.consume_cancel_grammar(): + for _, state, _ in grammar_manager.grammar_queue: + state.set_finish_with_abort("Aborted by pause", notify_client=True) + + # On resume, flush specs buffered while paused even when no new request + # arrives this iteration. This must run before the ``if not ready: + # return`` guard below, which would otherwise strand buffered specs + # until the next inbound request. Specs aborted while paused are reaped + # in place (terminating finish + state cleanup) rather than admitted, so + # they don't burn a scheduler slot or leak their rid — see + # ``_reap_or_keep_buffered_spec``. + if not self._pause.admit_blocked and self._pause.buffered_specs: + specs = [ + spec + for spec in self._pause.take_buffered_specs() + if self._reap_or_keep_buffered_spec(spec) + ] + if specs: + self.scheduler.submit_requests(specs) + + # Partition new requests by grammar readiness. Compile-bound requests + # are queued in GrammarManager and admitted in a later iteration when + # their futures resolve (see _drain_ready_grammar_requests below). + ready = [] + for spec, state, bootstrap in zip( + new_req_specs, new_req_states, bootstrap_infos + ): + # Requests pre-marked finished (e.g. invalid session ID aborted + # in RequestHandler) skip grammar compilation entirely — we'd + # just be wasting a compile slot on a response we're about to + # abort anyway, and the terminal response would be delayed by + # the compile/timeout window. + if state.finished: + ready.append((spec, state, bootstrap)) + continue + if grammar_manager.process_req_with_grammar(state): + ready.append((spec, state, bootstrap)) + else: + grammar_manager.add_to_queue(spec, state, bootstrap) + + # Drain any previously-queued requests whose grammar just finished + # compiling. With attn_tp > 1 this also drives the per-iter all_gather + # that keeps grammar admission in sync across ranks. + ready.extend(grammar_manager.get_ready_grammar_requests()) + + if not ready: + return + + admitted_specs = [] + for spec, state, bootstrap in ready: + # Grammar-aborted (invalid grammar, timed-out compile, or missing + # backend) requests must not enter the scheduler — they have no + # valid grammar to mask logits with, and we don't want to spend a + # prefill slot on a request that's already finished. Publish the + # finish_reason directly so the client still gets a response. + if state.finished: + self.output_processor.publish_finished_at_admission( + spec.request_id, state + ) + continue + + if isinstance(self.kv_transfer, DisaggDecodeExecutor): + state.computed_length = state.input_length + self.output_processor.register(spec.request_id, state) + is_epd = self.epd_admission is not None and self._is_epd_request(state) + # EPD: DEFER the P->D sender registration to admission (in + # EpdPrefillAdmission.drain, just before submit_requests). Registering + # it now -- while the request is staged and NOT yet in the C++ scheduler + # -- would let DisaggPrefillExecutor.generate_events poll the sender and + # emit a BootstrappedEvent that the scheduler's requests_.at(rid) THROWS + # on (no such request yet). Non-EPD requests register now (submitted this + # same call). + if self.kv_transfer is not None and not is_epd: + self.kv_transfer.register(spec.request_id, bootstrap) + + if self.memory_executor is not None: + hashes = calc_l3_query_hashes(self.scheduler, spec.tokens) + if hashes and len(hashes) > self.prefetch_threshold: + hit_pages = self.memory_executor.query_l3_pages(hashes) + logger.debug( + "[cache_op] L3 query: rid=%s hash_pages=%s hit_pages=%s threshold=%s", + spec.request_id, + len(hashes), + hit_pages, + self.prefetch_threshold, + ) + spec.rolling_hashes = hashes + spec.storage_hit_pages = hit_pages + # EPD prefill: hold a request whose images are encode-routed OUT of the + # scheduler until its per-image embeddings have been received (started + # here, polled in EpdPrefillAdmission.drain, which registers the P->D + # sender + submits once ready). It is output_processor-registered above; + # the sender registration + submission are both deferred. Non-EPD + # requests admit immediately as before. Rank-identical because `ready` is + # rank-synced (recv_reqs broadcast + grammar gather). + if is_epd: + self.epd_admission.stage( + spec.request_id, state.multimodal_inputs.mm_items + ) + self._epd_staged[spec.request_id] = (spec, state, bootstrap) + else: + admitted_specs.append(spec) + + # Pause gate: while paused, withhold new requests from the scheduler + # (running requests keep stepping); buffered specs are flushed on resume + # above, ahead of any newly-admitted ones, preserving FIFO order. + # + # TODO(pause-fifo): recv_reqs() drains the socket non-blocking, so a + # generate request that arrived *before* a pause control message can be + # coalesced into the same batch and reach here after the pause flipped + # admit_blocked. Such a pre-pause request is buffered as post-pause work + # instead of running (wait) / being aborted (abort). Correct handling + # needs the batch processed as an ordered stream that respects the + # control request's FIFO position. Tracked as a follow-up; until then we + # warn when the coalescing condition is observed so it is not silent. + if self._pause.admit_blocked: + if admitted_specs and not pause_blocked_before: + logger.warning( + "Pause engaged in the same recv batch as %d generate " + "request(s) (rids=%s); their FIFO order relative to the " + "pause is not preserved, so a pre-pause request may be " + "buffered as post-pause work and run only after resume. " + "See TODO(pause-fifo).", + len(admitted_specs), + [spec.request_id for spec in admitted_specs], + ) + self._pause.buffer_specs(admitted_specs) + return + + if admitted_specs: + self.scheduler.submit_requests(admitted_specs) + + @nvtx_range("loop:commit", color="rapids") + def _commit_forward_results( + self, + forward_op, + results: ModelExecutionResult, + on_first_token=None, + ): + self.request_handler.forward_ct += 1 + forward_mode = ForwardMode.from_num_extends( + forward_op.num_extends(), + len(forward_op.request_ids), + ) + self.request_handler._profile_batch_predicate(forward_mode) + + # post_process_forward_op calls sync() — after this, CPU tensors are ready + is_prefill_instance = isinstance(self.kv_transfer, DisaggPrefillExecutor) + request_changes = self.output_processor.post_process_forward_op( + forward_op, + results, + is_prefill_instance=is_prefill_instance, + on_first_token=on_first_token, + ) + + # Accumulate decode stats from synced results (no GPU sync) + if forward_op.num_extends() <= 0: + bs = len(forward_op.request_ids) + self.model_executor.accumulate_decode_stats(results, bs) + + return request_changes + + def _get_forward_op(self, execution_plan): + """Return the next forward op from the given plan, or None if there is nothing to run.""" + forward_ops = execution_plan.forward + if len(forward_ops) == 0 or len(forward_ops[0].request_ids) == 0: + return None + return forward_ops[0] + + def _handle_flat_oom_terminals(self, execution_plan) -> None: + """Surface flat-KV OOM terminals to their clients as abort finishes. + + The C++ flat scheduler terminalizes a request that can never fit the + flat pool (AbortEvent inside the scheduler; the reaper reclaims its + resources) and reports its id on ``plan.flat_oom_request_ids`` + (always empty on radix builds). The scheduler side is already fully + torn down — do NOT send a ForwardEvent.Abort back — but the client is + still waiting on the response stream, so finish the request with an + abort here (mirrors the PD FailedEvent handling above, minus the + scheduler abort). + """ + oom_rids = getattr(execution_plan, "flat_oom_request_ids", None) + if not oom_rids: + return + for rid in oom_rids: + state = self.output_processor.rid_to_state.get(rid) + if state is None: + # rid already gone (e.g. a client abort raced ahead). + logger.debug( + "flat OOM terminal for rid=%s: state missing; skipping", rid + ) + continue + if state.finished: + # Already carries a finish (an abort raced ahead of the + # terminal). C++ reports this rid exactly once and no future + # forward op will reap it, so resolve it here (same orphan + # rule as _reap_or_keep_buffered_spec). + self.output_processor.reap_finished_orphan(rid, state) + continue + state.set_finish_with_abort( + "flat KV cache cannot fit this request: prompt exceeds pool " + "capacity (OOM)" + ) + self.output_processor.publish_finished_at_admission(rid, state) + + def _process_kv_transfer_events(self, kv_transfer_events: list) -> list: + processed = [] + for event in kv_transfer_events: + processed.append(event) + if isinstance(event, PD.SucceededEvent) and isinstance( + self.kv_transfer, DisaggPrefillExecutor + ): + req_id = event.request_id + processed.extend(self.output_processor.finish_prefill_request(req_id)) + elif isinstance(event, PD.RemotePrefillDoneEvent): + req_id = event.request_id + bootstrap_token = event.bootstrap_token + self.output_processor.on_remote_prefill_done(req_id, bootstrap_token) + if isinstance(self.kv_transfer, DisaggDecodeExecutor): + candidate_info = self.kv_transfer.pop_remote_spec_candidate_ids( + req_id + ) + if candidate_info is not None: + req_pool_idx, candidate_ids = candidate_info + self.model_executor.write_remote_spec_candidate_ids( + req_pool_idx, candidate_ids + ) + elif isinstance(event, PD.FailedEvent): + # A PD/EPD transfer failed: the decode KV receiver timed out (e.g. the + # prefill aborted on embedding timeout so the KV never arrives), or a + # transfer errored. The C++ scheduler's FailedEvent handler is a no-op, + # so without this the request is never finished and the CLIENT HANGS + # FOREVER (the decode is its response stream). Finish it with an abort + # (streams the error to the client) and abort it in the scheduler so its + # slot/KV is freed. AbortEvent is valid from the decode-waiting state + # (forward_events.cpp: AbortEvent(Prefilling&&) -> Finished). + req_id = event.request_id + state = self.output_processor.rid_to_state.get(req_id) + if state is not None: + state.set_finish_with_abort( + "PD/EPD remote transfer failed or timed out" + ) + self.output_processor.publish_finished_at_admission(req_id, state) + abort = ForwardEvent.Abort() + abort.request_id = req_id + processed.append(abort) + + return processed + + def _get_load(self): + """Return load metrics for the DP load balancer.""" + from tokenspeed.runtime.engine.io_struct import GetLoadReqOutput + + available = self.scheduler.available_kv_pages() + num_total_pages = self.max_total_num_tokens // self.server_args.block_size + num_used_pages = num_total_pages - available + num_waiting = self.scheduler.waiting_size() + # num_reqs: running + waiting (used by SHORTEST_QUEUE balancing) + num_running = len(self.output_processor.rid_to_state) + return GetLoadReqOutput( + dp_rank=self.dp_rank, + num_reqs=num_running + num_waiting, + num_waiting_reqs=num_waiting, + num_pages=num_used_pages, + ) + + def _dp_sync_and_check(self, forward_op) -> DpForwardMetadata: + """Synchronize DP ranks with CPU-only metadata. + + All ranks call this before GPU forward work. The gathered metadata is + used for eager token-aware collectives and for choosing a common padded + CUDA graph shape during decode. + """ + import torch.distributed as dist + + executes_model_forward = _forward_op_executes_model_forward( + forward_op, + is_disagg_decode=isinstance(self.kv_transfer, DisaggDecodeExecutor), + ) + num_tokens = sum(forward_op.input_lengths) if executes_model_forward else 0 + batch_size = len(forward_op.request_ids) if executes_model_forward else 0 + if not executes_model_forward: + forward_mode = ForwardMode.IDLE + else: + forward_mode = ForwardMode.from_num_extends( + forward_op.num_extends(), + batch_size, + ) + + self._dp_local_info[0, 0] = num_tokens + self._dp_local_info[0, 1] = batch_size + self._dp_local_info[0, 2] = int(forward_mode) + dist.all_gather_into_tensor( + self._dp_global_info, + self._dp_local_info, + group=self.world_cpu_group, + ) + global_num_tokens = self._dp_global_info[:, 0].tolist() + global_batch_size = self._dp_global_info[:, 1].tolist() + global_forward_mode = self._dp_global_info[:, 2].tolist() + any_rank_has_work = max(global_num_tokens) > 0 + need_idle_forward = num_tokens == 0 and any_rank_has_work + all_decode_or_idle = all( + mode + in ( + int(ForwardMode.DECODE), + int(ForwardMode.IDLE), + ) + for mode in global_forward_mode + ) + # Replicated prefill-graph gate (see PrefillGraph._select_bucket). + all_extend = all( + mode == int(ForwardMode.EXTEND) for mode in global_forward_mode + ) + return DpForwardMetadata( + global_num_tokens=global_num_tokens, + global_batch_size=global_batch_size, + global_forward_mode=global_forward_mode, + all_decode_or_idle=all_decode_or_idle, + all_extend=all_extend, + need_idle_forward=need_idle_forward, + ) + + def _get_scheduler_stats(self): + """Query scheduler for page usage and queue depth.""" + available = self.scheduler.available_kv_pages() + active = self.scheduler.active_kv_pages() + num_total_pages = self.max_total_num_tokens // self.server_args.block_size + return { + "num_active_pages": active, + "num_cached_pages": num_total_pages - available, + "num_queue_reqs": self.scheduler.waiting_size(), + } + + def _record_scheduler_iteration_metrics( + self, stats: dict, num_iteration_tokens: int + ) -> None: + self.metrics.record_scheduler_iteration( + running=len(self.output_processor.rid_to_state), + waiting=stats["num_queue_reqs"], + num_active_pages=stats["num_active_pages"], + num_total_pages=self.max_total_num_tokens // self.server_args.block_size, + num_iteration_tokens=num_iteration_tokens, + ) + + # ------------------------------------------------------------------ + # Pause / resume helpers + # ------------------------------------------------------------------ + + def _reset_caches_for_release(self) -> None: + """Invalidate the prefix/radix cache before KV is discarded on release. + + KV pages are re-mapped + zeroed on wake, so any retained prefix entry + would be stale. The unsafe case (prefix caching on with no reset) is + rejected up front in ``MemoryOccupationController.handle_release`` via + ``kv_cache_release_allowed``, so by the time we get here either a reset + exists or prefix caching is off (nothing to invalidate). + """ + reset = getattr(self.scheduler, "reset_prefix_cache", None) + if callable(reset): + reset() + + def _kv_pools(self) -> list: + """All KV pools whose pages are tagged ``kv_cache`` — the target pool and + the draft pool in speculative-decoding runs. Release/repair must walk the + SAME set, so both derive it here rather than enumerating pools by hand.""" + pools = [] + for attr in ("token_to_kv_pool", "draft_token_to_kv_pool"): + pool = getattr(self.model_executor, attr, None) + if pool is not None: + pools.append(pool) + return pools + + def _kv_repair_after_wake(self) -> None: + """Zero re-mapped KV buffers (garbage after re-map) for every KV pool, + including the draft pool in spec-decode runs — its allocations are tagged + ``kv_cache`` too, so a wake that skipped it would feed the draft model + stale KV. FP8 KV scales ride with the weights region, so no scale reset + is needed here.""" + for pool in self._kv_pools(): + if hasattr(pool, "clear_kv_buffers"): + pool.clear_kv_buffers() + + def _paused_idle_step(self, prev_forward_op=None, prev_results=None) -> None: + """Run one iteration under ``PAUSED_ALL`` (keep mode): no new forward + work, but keep DP ranks in lockstep, service the drain check, and yield + the CPU so the freeze does not busy-spin a core.""" + if prev_results is not None: + request_changes = self._commit_forward_results( + prev_forward_op, prev_results + ) + advance_forward(self.scheduler, request_changes) + self._publish_scheduler_kv_events() + + if self.has_dp: + dp_metadata = self._dp_sync_and_check(None) + # While memory is released the weights region is unmapped; an idle + # forward runs the model and would read freed memory. All DP ranks + # release together, so skipping the idle forward stays consistent + # across ranks (the small DP sync above still runs to keep lockstep). + if dp_metadata.need_idle_forward and not self._pause.released: + self.model_executor.execute_idle_forward( + dp_metadata.global_num_tokens, + dp_metadata.global_batch_size, + dp_metadata.all_decode_or_idle, + ) + + self._pause.maybe_finish_drain(self.scheduler) + time.sleep(_PAUSED_IDLE_SLEEP_S) + + # ------------------------------------------------------------------ + # Event loops + # ------------------------------------------------------------------ + + def event_loop(self): + """Non-overlapping scheduler loop.""" + while True: + self._process_new_requests() + # EPD prefill: admit requests whose async embedding receives completed + # this cycle (rank-synced). Fixed position right after + # _process_new_requests so the drain's TP collective ordering is + # rank-identical every cycle. + self._drain_ready_epd_embeddings() + self._commit_cache_results() + if self._pause.forward_blocked: + self._paused_idle_step() + continue + execution_plan = self.scheduler.next_execution_plan() + self._publish_scheduler_kv_events() + self._handle_flat_oom_terminals(execution_plan) + self._submit_cache_ops(execution_plan) + + forward_op = self._get_forward_op(execution_plan) + self._flush_mamba_retract_states(forward_op) + + stats = self._get_scheduler_stats() + num_iter_tokens = ( + sum(forward_op.input_lengths) if forward_op is not None else 0 + ) + + # DP sync: all ranks must participate even when idle. + dp_metadata = None + if self.has_dp: + dp_metadata = self._dp_sync_and_check(forward_op) + if dp_metadata.need_idle_forward: + self.model_executor.execute_idle_forward( + dp_metadata.global_num_tokens, + dp_metadata.global_batch_size, + dp_metadata.all_decode_or_idle, + ) + self._record_scheduler_iteration_metrics(stats, num_iter_tokens) + continue + + request_changes = [] + + if forward_op is not None: + sampling_params_list = self._gather_sampling_params(forward_op) + grammar_inputs = self._gather_grammar_state(forward_op) + self._mark_stats_scheduled(forward_op) + results, on_first_token = self._dispatch_forward( + forward_op, + sampling_params_list, + execution_plan, + dp_metadata=dp_metadata, + stats=stats, + grammar_inputs=grammar_inputs, + ) + if results is not None: + request_changes.extend( + self._commit_forward_results( + forward_op, results, on_first_token + ) + ) + + if self.kv_transfer is not None: + kv_transfer_events = self.kv_transfer.generate_events() + request_changes.extend( + self._process_kv_transfer_events(kv_transfer_events) + ) + + if request_changes: + advance_forward(self.scheduler, request_changes) + self._publish_scheduler_kv_events() + + # Resolve a deferred abort/wait pause reply once in-flight work drains. + self._pause.maybe_finish_drain(self.scheduler) + + self._record_scheduler_iteration_metrics(stats, num_iter_tokens) + + def _mark_stats_scheduled(self, forward_op) -> None: + # Stamp the pre-forward "scheduled" time on each request's stats tracker + # so the queue/prefill split is anchored before the forward (idempotent: + # only the first forward a request appears in sets it). --enable-log-request-stats. + if not self.server_args.enable_log_request_stats or forward_op is None: + return + now = time.time() + rid_to_state = self.output_processor.rid_to_state + for rid in forward_op.request_ids: + st = rid_to_state.get(rid) + if st is not None: + st.stats.mark_scheduled(now) + + def _gather_sampling_params(self, forward_op) -> list[SamplingParams]: + """Look up per-request SamplingParams from the output processor. The + sampling backend does its own flip detection + RNG state management + internally, so we only need the scalar params here.""" + return [ + self.output_processor.rid_to_state[rid].sampling_params + for rid in forward_op.request_ids + ] + + def _gather_grammar_state(self, forward_op) -> GrammarStepInputs | None: + """Build ``GrammarStepInputs`` for the current batch, or ``None``. + + Returns ``None`` when no request in this batch has a grammar — the + model_executor short-circuits then. Otherwise carries the grammars + list + per-EXTEND-slot ``advance_mask`` (False on intermediate + chunked-prefill chunks, since the sampled token is discarded by + post_process and must not advance the matcher). + """ + rid_to_state = self.output_processor.rid_to_state + grammars = [rid_to_state[rid].grammar for rid in forward_op.request_ids] + if not any(grammars): + return None + + advance_mask = None + num_extends = forward_op.num_extends() + if num_extends > 0: + bs = len(forward_op.request_ids) + extend_prefix_lens = forward_op.extend_prefix_lens + extend_input_lengths = forward_op.input_lengths[:num_extends] + advance_mask = [True] * bs + for i in range(num_extends): + rid = forward_op.request_ids[i] + # This chunk completes prefill iff it processes the final + # token of the prompt; intermediate chunks don't. + advance_mask[i] = ( + extend_prefix_lens[i] + extend_input_lengths[i] + >= rid_to_state[rid].input_length + ) + + return GrammarStepInputs(grammars=grammars, advance_mask=advance_mask) + + def event_loop_overlap(self): + """ + Overlapping scheduler loop: post-process the previous step's results + while the current step's forward pass is in flight. + """ + # EPD invariant: the async embedding drain (EpdPrefillAdmission.drain) + # that admits EPD requests runs ONLY in event_loop(), never here. A + # prefill node that receives encode embeddings must therefore run the + # non-overlap loop -- should_use_overlap_schedule enforces this by forcing + # prefill -> non-overlap. Assert it rather than trusting that external + # coupling: if a prefill ever reached this loop, every EPD request would + # stage into the admission controller and hang forever with no drain. + assert self.epd_admission is None, ( + "EPD prefill must run the non-overlap event_loop(); the embedding " + "drain is not wired into event_loop_overlap()" + ) + prev_results: ModelExecutionResult = None + prev_forward_op = None + + while True: + # Order this iter's default-stream writes (KVAllocator, + # update_block_table, prefix_cache writes to req_to_page) + # after the prev iter's forward on execution_stream that + # reads the same tensor. Non-blocking on host. + torch.cuda.default_stream().wait_stream( + self.model_executor.execution_stream + ) + self._process_new_requests() + self._commit_cache_results() + if self._pause.forward_blocked: + # Freeze: commit any in-flight (overlapped) step — a forward + # already on the GPU can't be un-launched — then idle. + self._paused_idle_step(prev_forward_op, prev_results) + prev_results = None + prev_forward_op = None + continue + execution_plan = self.scheduler.next_execution_plan() + self._publish_scheduler_kv_events() + self._handle_flat_oom_terminals(execution_plan) + + self._submit_cache_ops(execution_plan) + + forward_op = self._get_forward_op(execution_plan) + self._flush_mamba_retract_states(forward_op) + + stats = self._get_scheduler_stats() + num_iter_tokens = ( + sum(forward_op.input_lengths) if forward_op is not None else 0 + ) + + grammar_inputs = None + if forward_op is not None: + # Gather both sampling params and grammar state BEFORE the + # prev_results commit below — that commit can finish requests + # and pop them from output_processor.rid_to_state, which would + # KeyError when we look up rids that are still in the current + # forward_op. + sampling_params_list = self._gather_sampling_params(forward_op) + grammar_inputs = self._gather_grammar_state(forward_op) + + # DP sync: all ranks must participate even when idle. + dp_metadata = None + if self.has_dp: + dp_metadata = self._dp_sync_and_check(forward_op) + if dp_metadata.need_idle_forward: + if prev_results is not None: + request_changes = self._commit_forward_results( + prev_forward_op, prev_results + ) + advance_forward(self.scheduler, request_changes) + self._publish_scheduler_kv_events() + prev_results = None + prev_forward_op = None + self.model_executor.execute_idle_forward( + dp_metadata.global_num_tokens, + dp_metadata.global_batch_size, + dp_metadata.all_decode_or_idle, + ) + self._record_scheduler_iteration_metrics(stats, num_iter_tokens) + continue + + # ---- dispatch current forward first (async GPU launch) ---- + # Issue curr's forward before committing prev so the GPU runs curr + # while the CPU syncs/post-processes prev. Committing prev first + # would block the CPU on prev's copy_event and leave the GPU idle + # until dispatch — visible as a gap between forwards in the trace. + # + # Eager grammar exception: setup_grammar_step reads each matcher's + # current state to fill the bitmask. Under the overlap pattern the + # matcher hasn't been advanced yet by prev's accept_token (commit + # below), so the fill would use a one-step-stale state and let the + # model sample a token the matcher then rejects. Capturable + # grammar dodges this with an in-graph hostfunc that advances + # before fill; eager has no equivalent, so we commit prev first + # whenever this batch carries grammars. Costs the dispatch/commit + # overlap for grammar batches but is correct. + request_changes = [] + curr_has_grammar = grammar_inputs is not None + eager_grammar_needs_advance = ( + curr_has_grammar + and prev_results is not None + and self.model_executor.eager_grammar_buffers is not None + ) + if eager_grammar_needs_advance: + request_changes.extend( + self._commit_forward_results(prev_forward_op, prev_results) + ) + prev_results = None + prev_forward_op = None + + curr_results = None + if forward_op is not None: + self._mark_stats_scheduled(forward_op) + curr_results, _ = self._dispatch_forward( + forward_op, + sampling_params_list, + execution_plan, + dp_metadata=dp_metadata, + stats=stats, + grammar_inputs=grammar_inputs, + ) + + # ---- post-process previous step (overlapped with current forward) ---- + if prev_results is not None: + request_changes.extend( + self._commit_forward_results(prev_forward_op, prev_results) + ) + + # ---- collect KV transfer events ---- + if self.kv_transfer is not None: + kv_transfer_events = self.kv_transfer.generate_events() + request_changes.extend( + self._process_kv_transfer_events(kv_transfer_events) + ) + + if request_changes: + advance_forward(self.scheduler, request_changes) + self._publish_scheduler_kv_events() + + # Resolve a deferred abort/wait pause reply once in-flight work drains. + self._pause.maybe_finish_drain(self.scheduler) + + self._record_scheduler_iteration_metrics(stats, num_iter_tokens) + + prev_results = curr_results + prev_forward_op = forward_op + + +def run_event_loop( + server_args: ServerArgs, + port_args: PortArgs, + pipe_writer, +): + mapping = server_args.mapping + gpu_id = mapping.rank % mapping.nprocs_per_node + server_args.base_gpu_id + attn_tp_rank = mapping.attn.tp_rank + dp_rank = mapping.attn.dp_rank + global_rank = mapping.rank + + setproctitle.setproctitle(f"tokenspeed::scheduler_{dp_rank}") + faulthandler.enable() + parent_process = psutil.Process().parent() + register_usr_signal() + + prefix = f" ATTN TP RANK {attn_tp_rank}" + configure_logger(server_args, prefix=prefix) + + try: + if server_args.disaggregation_mode == "encode": + # The encode role is LM-free; run the lightweight vision-tower loop + # instead of building the full EventLoop (KV/LM scheduler). + from tokenspeed.runtime.pd.epd.encode_loop import ( + run_encode_loop, + ) + + run_encode_loop(server_args, port_args, pipe_writer, gpu_id, global_rank) + return + + event_loop = EventLoop( + server_args, + port_args, + gpu_id, + attn_tp_rank, + dp_rank, + global_rank, + ) + pipe_writer.send( + { + "status": "ready", + "max_total_num_tokens": event_loop.max_total_num_tokens, + "max_req_input_len": event_loop.max_req_input_len, + "max_num_seqs": server_args.max_num_seqs, + "chunked_prefill_size": server_args.chunked_prefill_size, + "max_model_len": event_loop.model_config.context_len, + } + ) + + if event_loop.has_dp: + # All DP schedulers must finish initialization before any rank enters + # the loop and starts the first DP metadata collective. + dist.barrier(group=event_loop.world_cpu_group) + + if event_loop.use_overlap_schedule: + event_loop.event_loop_overlap() + else: + event_loop.event_loop() + + except Exception: + traceback = get_exception_traceback() + logger.error("Scheduler hit an exception: %s", traceback) + parent_process.send_signal(signal.SIGUSR1) diff --git a/python/tokenspeed/runtime/engine/exceptions.py b/python/tokenspeed/runtime/engine/exceptions.py new file mode 100644 index 0000000..8ef11ad --- /dev/null +++ b/python/tokenspeed/runtime/engine/exceptions.py @@ -0,0 +1,33 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Engine-path exceptions for the async frontend.""" + +from __future__ import annotations + + +class EngineGenerateError(ValueError): + """Raised when a single request errored out on the scheduler side. + + Subclasses :class:`ValueError` so existing ``except ValueError`` + blocks in the serving layer continue to catch it. The narrower + type is available for new code that wants to discriminate + generate-errors from request-shape validation errors. + """ diff --git a/python/tokenspeed/runtime/engine/generation_output_processor.py b/python/tokenspeed/runtime/engine/generation_output_processor.py new file mode 100644 index 0000000..69162c0 --- /dev/null +++ b/python/tokenspeed/runtime/engine/generation_output_processor.py @@ -0,0 +1,995 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import time +from typing import TYPE_CHECKING + +from tokenspeed.runtime.engine.io_struct import BatchTokenIDOut +from tokenspeed.runtime.engine.request_stats import ( + NOOP_STATS, + RequestStats, + RequestStatsTracker, +) +from tokenspeed.runtime.engine.request_types import ( + ABORT_CODE, + FINISH_ABORT, + FINISH_LENGTH, + FINISH_MATCHED_STR, + FINISH_MATCHED_TOKEN, + INIT_INCREMENTAL_DETOKENIZATION_OFFSET, + BaseFinishReason, +) +from tokenspeed.runtime.engine.scheduler_utils import ( + make_abort_event, + make_extend_result_event, + make_finish_event, + make_update_reserve_tokens_event, +) +from tokenspeed.runtime.sampling.sampling_params import SamplingParams + +if TYPE_CHECKING: + from tokenspeed.runtime.engine.io_struct import TokenizedGenerateReqInput + from tokenspeed.runtime.execution.types import ModelExecutionResult + from tokenspeed.runtime.metrics.collector import EngineMetrics + from tokenspeed.runtime.grammar.base_grammar_backend import ( + BaseGrammarObject, + ) + +from tokenspeed.runtime.utils import get_colorful_logger +from tokenspeed.runtime.utils.nvtx import nvtx_range + +logger = get_colorful_logger(__name__) + +DEFAULT_FORCE_STREAM_INTERVAL = 50 + + +class RequestState: + """Per-request state needed for incremental streaming output. + + Extracts only the fields required by process_output from the incoming + request. Does not hold a reference to Req or any scheduler object. + """ + + def __init__( + self, + prompt_input_ids: list[int], + sampling_params: SamplingParams, + stream: bool, + tokenizer, + eos_token_ids: list[int] = None, + return_logprob: bool = False, + top_logprobs_num: int = 0, + token_ids_logprob: list[int] | None = None, + multimodal_inputs=None, + prompt_input_ids_unpadded: list[int] | None = None, + created_time: float = 0.0, + ) -> None: + # --- Extracted from recv_req (immutable) --- + self.prompt_input_ids: list[int] = prompt_input_ids + self.prompt_input_ids_unpadded: list[int] = ( + prompt_input_ids_unpadded + if prompt_input_ids_unpadded is not None + else prompt_input_ids + ) + self.multimodal_inputs = multimodal_inputs + self.sampling_params = sampling_params + self.stream = stream + self.eos_token_ids = eos_token_ids + self.tokenizer = tokenizer + self.computed_length = 0 + self.return_logprob = return_logprob + self.top_logprobs_num = top_logprobs_num + self.token_ids_logprob = token_ids_logprob + + # --- generation state (updated with forward step) --- + self.output_ids: list[int] = [] + self.finished_reason: BaseFinishReason | None = None + self.cached_tokens: int = 0 + self.prefix_len: int = 0 + self.spec_verify_ct: int = 0 + self.accept_draft_tokens: float | None = None + + # request stats (host-side); tracker attached only with --enable-log-request-stats + self.created_time: float = created_time + self.stats: RequestStatsTracker = NOOP_STATS + # Sampled-token logprobs, accumulated per generated token. + # None when return_logprob is False. + self.output_token_logprobs_val: list[float] | None = ( + [] if return_logprob else None + ) + self.output_token_logprobs_idx: list[int] | None = ( + [] if return_logprob else None + ) + + # --- Streaming bookkeeping (internal) --- + self._surr_offset: int | None = None + self._read_offset: int | None = None + self.decoded_text: str = "" + self.send_token_offset: int = 0 + self.send_decode_id_offset: int = 0 + self.finished_output: bool = False + + # abort related + self.to_abort = False + self.to_abort_message = None + # Client-initiated aborts skip streaming a finish (the TM already tore + # down its state). Pause-initiated aborts set this so the passive client + # still receives a terminating finish. + self.abort_notify_client = False + + # cached tokenizer ids + self._eos_token_id_cached = None + self._additional_stop_token_ids_cached = None + + # Constrained-decoding state. + self.grammar: BaseGrammarObject | None = None + self.grammar_key: tuple[str, str] | None = None + self.grammar_queued_ts: float = 0.0 + + def set_finish_with_abort(self, message: str, notify_client: bool = False) -> None: + """Mark this request as aborted with ``message``; finished_reason is + materialized immediately so callers don't need a check_finished() pass. + + ``notify_client`` streams a terminating finish to the client (used for + pause-initiated aborts, where the client did not tear down its state). + """ + self.to_abort = True + self.to_abort_message = message + self.abort_notify_client = notify_client + self.finished_reason = FINISH_ABORT(message=message) + + @classmethod + def from_recv_req( + cls, + recv_req: TokenizedGenerateReqInput, + tokenizer, + eos_token_ids: list[int], + ) -> RequestState: + return cls( + prompt_input_ids=recv_req.input_ids, + sampling_params=recv_req.sampling_params, + stream=recv_req.stream, + tokenizer=tokenizer, + eos_token_ids=eos_token_ids, + return_logprob=getattr(recv_req, "return_logprob", False), + top_logprobs_num=getattr(recv_req, "top_logprobs_num", 0), + token_ids_logprob=getattr(recv_req, "token_ids_logprob", None), + multimodal_inputs=getattr(recv_req, "multimodal_inputs", None), + prompt_input_ids_unpadded=getattr(recv_req, "input_ids_unpadded", None), + created_time=getattr(recv_req, "created_time", 0.0), + ) + + @property + def finished(self) -> bool: + return self.finished_reason is not None + + @property + def input_length(self) -> int: + return len(self.prompt_input_ids) + + @property + def output_length(self) -> int: + return len(self.output_ids) + + @property + def prefill_finished(self): + return self.computed_length >= self.input_length + + def add_computed_length(self, incr: int): + self.computed_length += incr + + def maybe_extend_multimodal_mrope_positions(self) -> None: + mm = self.multimodal_inputs + if mm is None or mm.mrope_positions is None: + return + + target_len = self.input_length + self.output_length + current_len = mm.mrope_positions.shape[-1] + if current_len >= target_len: + return + + from tokenspeed.runtime.multimodal.mrope import ( + extend_mrope_positions_for_retracted_request, + ) + + mm.mrope_positions = extend_mrope_positions_for_retracted_request( + mm.mrope_positions, target_len - current_len + ) + + def release_pending_multimodal_features(self) -> None: + mm = self.multimodal_inputs + if mm is not None and hasattr(mm, "release_shm_features"): + mm.release_shm_features() + + def init_incremental_detokenize(self): + """Return (all_ids_from_surr_offset, read_offset_relative_to_surr).""" + if self._surr_offset is None or self._read_offset is None: + self._read_offset = len(self.prompt_input_ids_unpadded) + self._surr_offset = max( + self._read_offset - INIT_INCREMENTAL_DETOKENIZATION_OFFSET, 0 + ) + all_ids = self.prompt_input_ids_unpadded + self.output_ids + return ( + all_ids[self._surr_offset :], + self._read_offset - self._surr_offset, + ) + + def check_finished(self, skip_grammar_termination: bool = False): + + if self.finished: + return + + if self.to_abort: + self.finished_reason = FINISH_ABORT( + message=self.to_abort_message, + ) + return + + # When the capturable-grammar hostfunc is authoritative, the + # caller identifies the terminating token itself (see + # post_process_forward_op); firing here would re-trigger on + # every later token and trim content via trim_matched_stop. + if not skip_grammar_termination and self.grammar is not None: + if self.grammar.is_terminated(): + self.finished_reason = FINISH_MATCHED_TOKEN(matched=self.output_ids[-1]) + return + + if len(self.output_ids) >= self.sampling_params.max_new_tokens: + self.finished_reason = FINISH_LENGTH( + length=self.sampling_params.max_new_tokens + ) + return + + last_token_id = self.output_ids[-1] + + if not self.sampling_params.ignore_eos: + matched_eos = False + + # Check stop token ids + if self.sampling_params.stop_token_ids: + matched_eos = last_token_id in self.sampling_params.stop_token_ids + if self.eos_token_ids: + matched_eos |= last_token_id in self.eos_token_ids + if self._eos_token_id_cached is None: + self.set_cached_id() + if self._eos_token_id_cached is not None: + matched_eos |= last_token_id == self._eos_token_id_cached + if self._additional_stop_token_ids_cached: + matched_eos |= last_token_id in self._additional_stop_token_ids_cached + if matched_eos: + self.finished_reason = FINISH_MATCHED_TOKEN(matched=last_token_id) + return + + # Check stop strings + if len(self.sampling_params.stop_strs) > 0: + tail_str = self.tokenizer.decode( + self.output_ids[-(self.sampling_params.stop_str_max_len + 1) :] + ) + + for stop_str in self.sampling_params.stop_strs: + if stop_str in tail_str or stop_str in self.decoded_text: + self.finished_reason = FINISH_MATCHED_STR(matched=stop_str) + return + + def set_cached_id(self): + """Assign tokenizer and cache ids needed by check_finished().""" + eos_id = getattr(self.tokenizer, "eos_token_id", None) + self._eos_token_id_cached = int(eos_id) if eos_id is not None else None + extra = getattr(self.tokenizer, "additional_stop_token_ids", None) + self._additional_stop_token_ids_cached = ( + set(int(x) for x in extra) if extra else None + ) + + +class OutputProcesser: + """Streams generation output to the detokenizer. + + Logprob support is intentionally omitted. + """ + + # Upper bound on how long a pending abort stays buffered waiting for + # its matching register(). Generous — a client reorder of more than + # a few seconds is already pathological; 5 min gives plenty of slack + # while preventing unbounded growth on stray/post-completion aborts. + _PENDING_ABORT_TTL_S = 300.0 + + def __init__( + self, + send_to_tokenizer, + attn_tp_rank: int = 0, + spec_algorithm=None, + spec_num_tokens: int | None = None, + stream_interval: int = 1, + enable_log_request_stats: bool = False, + *, + metrics: EngineMetrics, + ) -> None: + # BatchTokenIDOut is pushed directly to + # ``send_to_tokenizer`` (AsyncLLM's input socket). The + # inline detokenizer inside AsyncLLM is the only + # detokenization path. + self.send_to_tokenizer = send_to_tokenizer + # Per-request logs fire on each DP replica's TP leader (attn_tp_rank == 0), + # NOT the global rank 0 — otherwise DP replicas > 0 would log nothing and + # their requests would be missing from the logs entirely. + self.attn_tp_rank = attn_tp_rank + self.spec_algorithm = spec_algorithm + self.spec_num_tokens = spec_num_tokens + self.stream_interval = stream_interval + self.metrics = metrics + self.enable_log_request_stats = enable_log_request_stats + # previous forward step ts, for host-side preempt timing + self._last_step_ts: float = 0.0 + self.log_cnt = 0 + self.rid_to_state: dict[str, RequestState] = {} + # rid → monotonic ts at which the abort was seen. Covers the + # "abort arrives before register()" race (pre-arrival reorder), + # plus grammar-queued aborts that publish_finished_at_admission + # handles. Entries for rids that never register are swept by TTL + # to keep this bounded across a long-running server. + self.pending_aborts: dict[str, float] = {} + + def log_accept_length(self, rid, request_state: RequestState): + # When --enable-log-request-stats is on, the richer RequestStats line (which + # already carries acc_len) replaces this one — see _log_request_stats. + if self.attn_tp_rank == 0 and not self.enable_log_request_stats: + logger.info( + "Req: %s Finish! Accept_num_tokens_avg: %s", + rid, + request_state.accept_draft_tokens, + ) + + def _log_request_stats( + self, rid: str, rs: RequestState, finish_time: float + ) -> None: + # Single guard for the whole stats path: no tracker (flag off) or non-zero + # rank => nothing to do. Keeps the forward-loop call sites trivial and the + # derivation in from_state total (it always sees a tracker). + if rs.stats is NOOP_STATS or self.attn_tp_rank != 0: + return + rs.stats.mark_finish(finish_time) + stats = RequestStats.from_state(rs, self.spec_algorithm, self.spec_num_tokens) + # Fused into the scheduler's per-request finish line (supersedes the + # Accept_num_tokens_avg variant in log_accept_length). + logger.info("Req: %s Finish! %s", rid, stats) + + def sweep_pending_aborts(self) -> None: + """Drop TTL-expired entries from ``pending_aborts``. + + Safe to call anytime. pending_aborts is insertion-ordered so we + can stop at the first non-expired entry. Called both inside + ``mark_abort`` (so adds are bounded) and periodically from the + event loop (so entries also age out when aborts stop arriving). + """ + cutoff = time.monotonic() - self._PENDING_ABORT_TTL_S + while self.pending_aborts: + oldest_rid = next(iter(self.pending_aborts)) + if self.pending_aborts[oldest_rid] >= cutoff: + break + self.pending_aborts.pop(oldest_rid) + + def mark_abort(self, rid: str, notify_client: bool = False): + """Mark a request for abort. Safe to call before or after register(). + + Routes through ``RequestState.set_finish_with_abort`` so + ``finished_reason`` is materialized immediately. Without that, + the gate ``request_state.to_abort and request_state.finished`` + in ``post_process_forward_op`` never fires (``.finished`` is + ``finished_reason is not None``), so the scheduler keeps + running the request until natural ``max_tokens``/EOS — the + cancelled request burns up to ``max_tokens`` forward steps and + latches a ``--max-num-seqs`` slot in the meantime. + + ``notify_client`` streams a terminating finish to the client (for + pause-initiated aborts; client-initiated aborts leave it False since + the tokenizer manager has already cleaned up its own state). + """ + state = self.rid_to_state.get(rid) + if state is not None: + msg = "Aborted by pause" if notify_client else "AbortReq from client" + state.set_finish_with_abort(msg, notify_client=notify_client) + return + + self.sweep_pending_aborts() + self.pending_aborts[rid] = time.monotonic() + + def register(self, rid, state): + self.rid_to_state[rid] = state + if self.enable_log_request_stats: + state.stats = RequestStatsTracker() + if self.pending_aborts.pop(rid, None) is not None: + # Same reasoning as ``mark_abort``: drive the abort all the + # way to ``finished_reason`` so the slot-release gate fires. + state.set_finish_with_abort("AbortReq from client") + + def publish_finished_at_admission(self, rid: str, state: RequestState) -> None: + """Stream a finish for a request that was finished before admission. + + Used for grammar-aborted requests (invalid/timed-out compile, missing + backend) so the client gets a finish_reason without us wasting a + scheduler slot or a forward step on them. + """ + self.rid_to_state[rid] = state + try: + state.finished_output = False + self.stream_output([rid], [state]) + finally: + state.release_pending_multimodal_features() + self.rid_to_state.pop(rid, None) + # This path replaces register() for grammar-aborted rids — + # drop any queued abort marker so pending_aborts doesn't leak + # and a reused rid isn't instantly re-aborted on next register. + self.pending_aborts.pop(rid, None) + + def reap_finished_orphan(self, rid: str, state: RequestState) -> None: + """Resolve a finished request that no future forward op will reap. + + Stream the terminating finish to a passive client (pause-initiated + aborts still have the client waiting on the stream); client-initiated + aborts already tore down their own state, so just drop the registered + state so the rid does not leak. + """ + if state.abort_notify_client: + self.publish_finished_at_admission(rid, state) + else: + self.rid_to_state.pop(rid, None) + + def _host_advance_matcher(self, completion, model_execution_results): + """Host-side fallback for the grammar matcher advance. + + Reads already-synced CPU tensors. Fires when no next step arrives to run + the hostfunc (e.g., last live request finished). + """ + grammars = completion.grammars or [] + stride = completion.tokens_per_req + bs = completion.bs + advance_mask = completion.advance_mask or [True] * bs + output_tokens = model_execution_results.output_tokens + accept_lengths = model_execution_results.output_lengths + terminated_at = [-1] * bs + for i, grammar in enumerate(grammars): + if ( + grammar is None + or grammar.finished + or grammar.is_terminated() + or not advance_mask[i] + ): + continue + n_accepted = int(accept_lengths[i].item()) + for j in range(n_accepted): + tok = int(output_tokens[i * stride + j].item()) + try: + grammar.accept_token(tok) + except Exception: + break + if grammar.is_terminated(): + terminated_at[i] = j + break + completion.terminated_at = terminated_at + + def add_computed_length(self, rids, input_lengths, extend_prefix_lens): + for i, rid in enumerate(rids): + if rid not in self.rid_to_state: + continue + if i < len(extend_prefix_lens): + self.rid_to_state[rid].computed_length = ( + input_lengths[i] + extend_prefix_lens[i] + ) # Avoid accumulation here so chunked prefill does not distort the value. + else: + self.rid_to_state[rid].add_computed_length(input_lengths[i]) + + @staticmethod + def _aggregate_spec_decode_step( + *, + forward_op, + output_lengths, + rid_to_state, + ) -> tuple[int, int]: + n_ext = forward_op.num_extends() + accepted = 0 + num_slots = 0 + for i in range(n_ext, len(forward_op.request_ids)): + rid = forward_op.request_ids[i] + rs = rid_to_state.get(rid) + if rs is None or not rs.prefill_finished: + continue + out_len = int(output_lengths[i].item()) + accepted += max(0, out_len - 1) + num_slots += 1 + return num_slots, accepted + + def _emit_spec_decode_metrics( + self, forward_op, model_execution_results: ModelExecutionResult + ) -> None: + if not self.metrics.enabled: + return + if forward_op.num_extends() > 0: + return + if self.spec_algorithm is None or self.spec_num_tokens is None: + return + if model_execution_results.output_lengths is None: + return + num_slots, accepted_draft_tokens = self._aggregate_spec_decode_step( + forward_op=forward_op, + output_lengths=model_execution_results.output_lengths, + rid_to_state=self.rid_to_state, + ) + if num_slots > 0: + self.metrics.record_spec_decode_step( + num_decode_slots=num_slots, + accepted_draft_tokens=accepted_draft_tokens, + draft_width=self.spec_num_tokens, + ) + + def add_cached_tokens(self, rids: list[str], extend_prefix_lens: list[int]) -> None: + for rid, prefix_len in zip(rids, extend_prefix_lens): + if rs := self.rid_to_state.get(rid): + rs.cached_tokens += max(0, prefix_len - rs.computed_length) + + def post_process_forward_op( + self, + forward_op, + model_execution_results: ModelExecutionResult, + is_prefill_instance: bool = False, + on_first_token=None, + ): + self.add_cached_tokens( + forward_op.request_ids, + forward_op.extend_prefix_lens, + ) + with nvtx_range("commit:sync", color="red"): + model_execution_results.sync() + + self._emit_spec_decode_metrics(forward_op, model_execution_results) + + # Wait briefly for the next step's build hostfunc to advance + # the matcher; if it doesn't come, advance on host. The lock + # on the completion ensures exactly one path wins. + grammar_completion = model_execution_results.grammar_completion + grammar_terminated_at = None + if grammar_completion is not None: + if not grammar_completion.event.wait(timeout=0.005): + with grammar_completion.lock: + if not grammar_completion.event.is_set(): + self._host_advance_matcher( + grammar_completion, model_execution_results + ) + grammar_completion.event.set() + grammar_terminated_at = grammar_completion.terminated_at + self.log_cnt += 1 + self.add_computed_length( + forward_op.request_ids, + forward_op.input_lengths, + forward_op.extend_prefix_lens, + ) + num_extends = forward_op.num_extends() + + # per-request stats timing (host-side, only when --enable-log-request-stats) + stats_now = time.time() if self.enable_log_request_stats else 0.0 + step_dt = 0.0 + prefilling_others = False + if self.enable_log_request_stats: + step_dt = ( + stats_now - self._last_step_ts if self._last_step_ts > 0.0 else 0.0 + ) + self._last_step_ts = stats_now + prefilling_others = num_extends > 0 + + request_changes = [] + stream_out_rids = [] + stream_out_states = [] + output_logprobs_list = ( + model_execution_results.output_logprobs.tolist() + if model_execution_results.output_logprobs is not None + else None + ) + # NaN-guard flags, aligned with forward_op.request_ids (None when disabled). + nan_flags_list = ( + model_execution_results.output_nan_flags.tolist() + if model_execution_results.output_nan_flags is not None + else None + ) + # Per-slot total prefill length as the OP sees it (C++ Request::PrefillSize()). + # After a flat retract the victim's generated tokens are rebased into the + # prefill window (RebasePrefill), so this can exceed the original prompt + # length that RequestState.prefill_finished compares against. + prefill_lengths = getattr(forward_op, "prefill_lengths", None) + pt = 0 + for i, rid in enumerate(forward_op.request_ids): + output_length = model_execution_results.output_lengths[i].item() + model_output_ids = model_execution_results.output_tokens.tolist()[ + pt : pt + output_length + ] + model_output_logprobs = ( + output_logprobs_list[pt : pt + output_length] + if output_logprobs_list is not None + else None + ) + is_decode_slot = i >= num_extends + if self.spec_num_tokens is not None and is_decode_slot: + pt += self.spec_num_tokens + else: + pt += output_length + + if rid not in self.rid_to_state: + # means it's delayed token, do not process + continue + + request_state: RequestState = self.rid_to_state[rid] + # scheduled_time is stamped pre-forward in the event loop (queue end) + + # Mid-chunk extend slot by the op's own prefill_lengths (rebased after + # flat retract; C++ owes no result and the sampled token is garbage). + # Fresh requests: prefill_length == prompt length, same as the gate below. + if ( + not is_decode_slot + and prefill_lengths is not None + and forward_op.extend_prefix_lens[i] + forward_op.input_lengths[i] + < prefill_lengths[i] + ): + continue + + # Do not output chunking result + if not request_state.prefill_finished: + continue + + request_state.stats.mark_prefill_done(stats_now) + if i >= num_extends: + request_state.stats.record_decode_step(step_dt, prefilling_others) + + nan_detected = nan_flags_list is not None and nan_flags_list[i] + if nan_detected and not request_state.finished: + request_state.finished_reason = FINISH_ABORT( + message=( + "Request terminated: numerical corruption (NaN logits" + " or out-of-vocab sample) detected during generation." + ), + err_type=ABORT_CODE.NumericalError, + ) + # Keep one sanitized token so accounting matches a mid-step finish. + model_output_ids = model_output_ids[:1] + if model_output_logprobs is not None: + model_output_logprobs = model_output_logprobs[:1] + self.metrics.record_nan_abort() + if self.attn_tp_rank == 0: + logger.warning( + "Req %s terminated: NaN detected in logits (or an" + " out-of-vocab sample escaped the sampler);" + " isolating it from the batch.", + rid, + ) + + # Notify caller of first output token (used by prefill node to hand off + # bootstrap token and speculative candidates to the KV transfer layer). + # NaN-terminated requests skip the handoff: their KV is suspect. + if on_first_token is not None and model_output_ids and not nan_detected: + bootstrap_token = int(model_output_ids[0]) + spec_candidate_ids = None + if model_execution_results.next_input_ids is not None and i < len( + model_execution_results.next_input_ids + ): + spec_candidate_ids = [ + int(x) + for x in model_execution_results.next_input_ids[i].tolist() + ] + + on_first_token( + rid, + forward_op.request_pool_indices[i], + bootstrap_token, + spec_candidate_ids, + ) + + if is_decode_slot and self.spec_algorithm is not None: + request_state.spec_verify_ct += 1 + + # With the capturable grammar pipeline the matcher is + # advanced by the hostfunc; here we just read which token + # (if any) terminated it so FINISH_MATCHED_TOKEN fires on + # the right token and check_finished skips the now-stale + # grammar.is_terminated() probe. + use_hostfunc = grammar_terminated_at is not None + advance_grammar = not use_hostfunc and request_state.grammar is not None + term_idx = ( + grammar_terminated_at[i] + if use_hostfunc and request_state.grammar is not None + else -1 + ) + new_ids = [] + for j, model_output_id in enumerate(model_output_ids): + request_state.output_ids.append(model_output_id) + if advance_grammar: + request_state.grammar.accept_token(model_output_id) + if ( + request_state.return_logprob + and request_state.output_token_logprobs_val is not None + and model_output_logprobs is not None + ): + request_state.output_token_logprobs_val.append( + model_output_logprobs[j] + ) + request_state.output_token_logprobs_idx.append(model_output_id) + if term_idx == j: + # Grammar termination takes precedence over + # length/EOS/stop_str at the same step (matching + # check_finished's original order). + request_state.finished_reason = FINISH_MATCHED_TOKEN( + matched=model_output_id + ) + else: + request_state.check_finished(skip_grammar_termination=use_hostfunc) + new_ids.append(model_output_id) + if request_state.finished: + request_state.accept_draft_tokens = ( + (len(request_state.output_ids) - 1) + / request_state.spec_verify_ct + if request_state.spec_verify_ct > 0 + else 0 + ) + self.log_accept_length(rid, request_state) + break + + # first output token == TTFT anchor + if request_state.output_ids: + request_state.stats.mark_first_token(stats_now) + + # For aborted requests, skip output to detokenizer (the tokenizer + # manager already cleaned up), just notify the scheduler to finish. + # Exception: pause-initiated aborts (abort_notify_client) leave a + # passive client that still needs a terminating finish streamed. + if request_state.to_abort and request_state.finished: + request_changes.append(make_extend_result_event(rid, new_ids)) + request_changes.append(make_finish_event(rid)) + if request_state.abort_notify_client: + stream_out_rids.append(rid) + stream_out_states.append(request_state) + self._log_request_stats(rid, request_state, stats_now) + request_state.release_pending_multimodal_features() + self.rid_to_state.pop(rid) + continue + + request_changes.append(make_extend_result_event(rid, new_ids)) + if is_prefill_instance: + # Prefill instances: never stream intermediate output to detokenizer. + # The finish packet is sent exactly once by finish_prefill_request() + # when SucceededEvent arrives (KV transfer complete). Sending output + # here would either give the client partial data or trigger a double- + # finish on the TM side. + pass + elif request_state.finished: + stream_out_rids.append(rid) + stream_out_states.append(request_state) + # Abort (vs Finish) keeps corrupted KV out of the prefix caches. + request_changes.append( + make_abort_event(rid) if nan_detected else make_finish_event(rid) + ) + self._log_request_stats(rid, request_state, stats_now) + request_state.release_pending_multimodal_features() + self.rid_to_state.pop(rid) + else: + stream_out_rids.append(rid) + stream_out_states.append(request_state) + if is_decode_slot: + request_changes.append( + make_update_reserve_tokens_event(rid, output_length) + ) + + self.stream_output(stream_out_rids, stream_out_states) + return request_changes + + def on_remote_prefill_done(self, req_id: str, bootstrap_token: int) -> None: + """Record the bootstrap token on a decode-node request (RemotePrefillDoneEvent). + + The bootstrap_token is the first real output token produced by the prefill node. + It is appended to output_ids so the decode side starts generation from the + correct position. + + bootstrap_token == -1 means the prefill side did not (or could not) supply a + token (e.g. it was generated on a rank whose ZMQ message arrived after the + success barrier had already been satisfied). + """ + if req_id not in self.rid_to_state: + return + if bootstrap_token == -1: + logger.warning( + "[on_remote_prefill_done] rid=%s received bootstrap_token=-1, skipping append to output_ids", + req_id, + ) + return + state = self.rid_to_state[req_id] + state.output_ids.append(bootstrap_token) + state.check_finished() + + def finish_prefill_request(self, req_id: str) -> list: + """Finish a prefill-instance request when KV transfer succeeds (SucceededEvent). + + Called by event_loop._process_kv_transfer_events at the correct moment — the + SucceededEvent itself drives the C++ FSM transition Decoding → Finished, + so we must NOT emit an additional make_finish_event here. + + We send a finished BatchTokenIDOut to the detokenizer so the Prefill TM + can resolve its HTTP coroutine and let the HTTP load balancer unblock. + Without this, the load balancer waits forever for the prefill side's HTTP + response while the decode side has already finished — client hangs + indefinitely. + """ + if req_id not in self.rid_to_state: + return [] + rs = self.rid_to_state.pop(req_id) + rs.release_pending_multimodal_features() + + # Ensure a finish reason is set so TokenizerManager marks the request done. + if not rs.finished: + rs.finished_reason = FINISH_LENGTH(length=len(rs.output_ids)) + rs.finished_output = False + # PD prefill node's terminal path (the other finish/abort logging lives in + # post_process_forward_op). Self-guarded, so a no-op when the flag is off. + self._log_request_stats(req_id, rs, time.time()) + self.stream_output([req_id], [rs]) + # SucceededEvent already finishes the C++ FSM; no extra FinishEvent needed + return [] + + def stream_output( + self, stream_out_rids: list[str], output_states: list[RequestState] + ) -> None: + """Collect per-step results and forward them to the detokenizer.""" + if len(output_states) == 0: + return + + rids_to_send = [] + finished_reasons = [] + decoded_texts: list[str] = [] + decode_ids_list = [] + read_offsets: list[int] = [] + output_ids = [] + output_multi_ids = [] + skip_special_tokens: list[bool] = [] + spaces_between_special_tokens: list[bool] = [] + no_stop_trim: list[bool] = [] + prompt_tokens: list[int] = [] + completion_tokens: list[int] = [] + cached_tokens: list[int] = [] + spec_verify_ct: list[int] = [] + batch_accept_draft_tokens: list[float] = [] + output_extra_infos: list[dict] = [] + output_token_logprobs_val: list[list[float]] = [] + output_token_logprobs_idx: list[list[int]] = [] + + for i, rs in enumerate(output_states): + # For finished requests, always output (unless already output) + if rs.finished: + if rs.finished_output: + # With the overlap schedule, a request will try to output twice and hit this line twice + # because of the one additional delayed token. This "continue" prevented the dummy output. + continue + rs.finished_output = True + should_output = True + else: + # For ongoing requests, use stream interval logic + if rs.stream: + stream_interval = getattr( + rs.sampling_params, "stream_interval", None + ) + if stream_interval is None: + stream_interval = self.stream_interval + should_output = ( + rs.output_length % stream_interval == 1 + if stream_interval > 1 + else rs.output_length % stream_interval == 0 + ) + else: + stream_interval = DEFAULT_FORCE_STREAM_INTERVAL + should_output = ( + rs.output_length == 1 or rs.output_length % stream_interval == 0 + ) + + if not should_output: + continue + + rids_to_send.append(stream_out_rids[i]) + send_token_offset = rs.send_token_offset + + finished_reasons.append( + rs.finished_reason.to_json() if rs.finished_reason else None + ) + decoded_texts.append(rs.decoded_text) + + decode_ids, read_offset = rs.init_incremental_detokenize() + decode_ids_list.append(decode_ids[rs.send_decode_id_offset :]) + rs.send_decode_id_offset = len(decode_ids) + + read_offsets.append(read_offset) + output_ids.append(rs.output_ids[send_token_offset:]) + rs.send_token_offset = rs.output_length + + output_multi_ids.append([]) + + skip_special_tokens.append(rs.sampling_params.skip_special_tokens) + spaces_between_special_tokens.append( + rs.sampling_params.spaces_between_special_tokens + ) + no_stop_trim.append(rs.sampling_params.no_stop_trim) + prompt_tokens.append(rs.input_length) + completion_tokens.append(rs.output_length) + cached_tokens.append(rs.cached_tokens) + + if self.spec_algorithm is not None: + spec_verify_ct.append(rs.spec_verify_ct) + batch_accept_draft_tokens.append(rs.accept_draft_tokens) + + output_extra_infos.append({"decode_prefix_len": rs.prefix_len}) + + if rs.return_logprob and rs.output_token_logprobs_val is not None: + # Send only the slice not yet shipped; send_token_offset was + # just advanced above, so use the logprob list tail. + n_new = rs.output_length - send_token_offset + output_token_logprobs_val.append( + rs.output_token_logprobs_val[-n_new:] if n_new > 0 else [] + ) + output_token_logprobs_idx.append( + rs.output_token_logprobs_idx[-n_new:] if n_new > 0 else [] + ) + else: + output_token_logprobs_val.append([]) + output_token_logprobs_idx.append([]) + + # Don't send empty batch to detokenizer + if len(rids_to_send) == 0: + return + + batch_id_out = BatchTokenIDOut( + rids=rids_to_send, + finished_reasons=finished_reasons, + decoded_texts=decoded_texts, + decode_ids=decode_ids_list, + read_offsets=read_offsets, + output_ids=output_ids, + output_multi_ids=output_multi_ids, + skip_special_tokens=skip_special_tokens, + spaces_between_special_tokens=spaces_between_special_tokens, + no_stop_trim=no_stop_trim, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + cached_tokens=cached_tokens, + spec_verify_ct=spec_verify_ct, + input_token_logprobs_val=[], + input_token_logprobs_idx=[], + output_token_logprobs_val=output_token_logprobs_val, + output_token_logprobs_idx=output_token_logprobs_idx, + input_top_logprobs_val=[], + input_top_logprobs_idx=[], + output_top_logprobs_val=[], + output_top_logprobs_idx=[], + input_token_ids_logprobs_val=[], + input_token_ids_logprobs_idx=[], + output_token_ids_logprobs_val=[], + output_token_ids_logprobs_idx=[], + output_hidden_states=[], + batch_accept_draft_tokens=batch_accept_draft_tokens, + output_extra_infos=output_extra_infos, + generated_time=time.time(), + ) + + # Push BatchTokenIDOut directly to AsyncLLM via the shared + # tokenizer-ipc socket. AsyncLLM runs IncrementalDetokenizer + # inline — there is no detokenizer subprocess anymore. + self.send_to_tokenizer.send_pyobj(batch_id_out) diff --git a/python/tokenspeed/runtime/engine/input_processor.py b/python/tokenspeed/runtime/engine/input_processor.py new file mode 100644 index 0000000..5182d69 --- /dev/null +++ b/python/tokenspeed/runtime/engine/input_processor.py @@ -0,0 +1,284 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Request tokenization helpers for the async frontend.""" + +from __future__ import annotations + +import asyncio +import json +import time +from typing import TYPE_CHECKING + +from tokenspeed.runtime.engine.io_struct import ( + EmbeddingReqInput, + GenerateReqInput, + SessionParams, + TokenizedEmbeddingReqInput, + TokenizedGenerateReqInput, +) +from tokenspeed.runtime.grammar.reasoning_structural_tag import ( + structural_tag_for_reasoning_json_schema, +) +from tokenspeed.runtime.multimodal.embedder import pad_input_tokens +from tokenspeed.runtime.multimodal.mrope import compute_mrope_positions +from tokenspeed.runtime.sampling.sampling_params import SamplingParams + +if TYPE_CHECKING: + from tokenspeed.runtime.engine.async_llm import AsyncLLM + + +class InputProcessor: + """Owns request-input logic: validation, tokenization, and the + tokenized-object prep for parallel-sampling fan-out. Callers + (``AsyncLLM``) stay thin — they route requests through this + class and then dispatch the resulting tokenized payloads to the + scheduler. + """ + + def __init__(self, engine: AsyncLLM): + self.engine = engine + + def _maybe_wrap_json_schema_for_reasoning(self, sampling: dict) -> None: + # Without this, xgrammar locks onto ``{`` at token 0 and the + # model can't emit ``…`` before the JSON. + if "json_schema" not in sampling: + return + reasoning_parser = getattr(self.engine.server_args, "reasoning_parser", None) + if not reasoning_parser: + return + try: + schema = sampling["json_schema"] + if isinstance(schema, str): + schema = json.loads(schema) + wrapped = structural_tag_for_reasoning_json_schema(reasoning_parser, schema) + except Exception as exc: + self.engine.logger.warning( + "reasoning-parser=%s: failed to wrap json_schema (%s); " + "falling back.", + reasoning_parser, + exc, + ) + return + if wrapped is None: + return + sampling.pop("json_schema", None) + sampling["structural_tag"] = wrapped + + def validate_request(self, obj: GenerateReqInput | EmbeddingReqInput) -> None: + """Reject cross-type requests before any other processing. + + An ``EmbeddingReqInput`` arriving at a generation-only engine + is a configuration mistake, not a runtime condition, so we + raise eagerly instead of letting it reach tokenization. + """ + if isinstance(obj, EmbeddingReqInput) and self.engine.is_generation: + raise ValueError("Embedding and rerank model requests are not supported.") + + async def tokenize_batch( + self, + objs: list[GenerateReqInput | EmbeddingReqInput], + ) -> list[TokenizedGenerateReqInput | TokenizedEmbeddingReqInput]: + """Tokenize a list of requests in parallel. + + Used by the batched fan-out path in ``AsyncLLM._handle_batch_request``. + The single-request path stays on ``tokenize_one_request`` — + avoiding the ``asyncio.gather`` hop keeps the hot path flat. + """ + return await asyncio.gather(*(self.tokenize_one_request(obj) for obj in objs)) + + async def tokenize_one_request( + self, + obj: GenerateReqInput | EmbeddingReqInput, + ) -> TokenizedGenerateReqInput | TokenizedEmbeddingReqInput: + """Tokenize one request without changing current behavior.""" + input_embeds = None + multimodal_inputs = None + input_ids_unpadded = None + input_text = obj.text + input_ids = obj.input_ids + + if obj.input_embeds is not None: + if self.engine.server_args.enable_prefix_caching: + raise ValueError( + "input_embeds is provided while prefix caching is enabled. " + "Please add `--no-enable-prefix-caching` when you launch the server " + "if you want to use input_embeds as inputs." + ) + input_embeds = obj.input_embeds + elif input_ids is None: + if self.engine.tokenizer is None: + raise ValueError( + "The engine initialized with skip_tokenizer_init=True cannot " + "accept text prompts. Please provide input_ids or re-initialize " + "the engine with skip_tokenizer_init=False." + ) + input_ids = self.engine.tokenizer.encode(input_text) + + precomputed_mm = ( + isinstance(obj, GenerateReqInput) + and obj.precomputed_multimodal_inputs is not None + ) + if precomputed_mm: + # Gateway-side preprocess path (e.g. SMG): mm tensors are already + # built by an upstream preprocessor and the input_ids carry the + # expanded placeholder tokens (im_token_id) at the right offsets. + # We still need to run pad_input_tokens so the engine's + # MultimodalEmbedder can plan encoder-token scatter ranges from each + # item's offsets — the bare placeholder token alone would not + # encode per-item uniqueness needed by the radix prefix layer. + if not self.engine.model_config.is_multimodal_active: + raise ValueError( + "precomputed_multimodal_inputs is provided for a text-only model." + ) + multimodal_inputs = obj.precomputed_multimodal_inputs + multimodal_inputs.ensure_pad_values() + # MRoPE-aware models (Qwen2/3-VL, …) require 3-axis position_ids + # derived from image_grid_thw + the image_token_id placeholders in + # input_ids. SMG ships precomputed mm inputs with mrope_* unset; if + # left None, model_executor falls back to a 1-D linear position + # override — silently degrading OCR accuracy. Compute them here, on + # the un-padded input_ids (so get_rope_index can still locate the + # image regions) BEFORE pad_input_tokens substitutes per-image + # pad_value over the placeholders, then pad for the embed splice. + if ( + input_ids is not None + and getattr(multimodal_inputs, "mrope_positions", None) is None + ): + mrope_positions, mrope_position_delta = compute_mrope_positions( + self.engine.model_config.hf_config, + list(input_ids), + multimodal_inputs.mm_items, + ) + multimodal_inputs.mrope_positions = mrope_positions + multimodal_inputs.mrope_position_delta = mrope_position_delta + if mrope_position_delta is not None: + multimodal_inputs.mrope_position_delta_scalar = int( + mrope_position_delta.flatten()[0].item() + ) + if input_ids is not None: + input_ids_unpadded = list(input_ids) + input_ids = pad_input_tokens(list(input_ids), multimodal_inputs) + + if self.engine.is_generation: + session_params = ( + SessionParams(**obj.session_params) if obj.session_params else None + ) + + input_token_num = len(input_ids) if input_ids is not None else 0 + if input_token_num >= self.engine.context_len: + raise ValueError( + f"The input ({input_token_num} tokens) is longer than the " + f"model's context length ({self.engine.context_len} tokens)." + ) + + max_new_tokens = obj.sampling_params.get("max_new_tokens") + # Resolve to a finite cap bounded by remaining context. Both + # Req.check_finished and RequestState.check_finished read this field; + # leaving it None lets a request reach the per-request page-table cap. + adjusted_max_new_tokens = self.engine.context_len - input_token_num + if max_new_tokens is None: + obj.sampling_params.update({"max_new_tokens": adjusted_max_new_tokens}) + elif max_new_tokens + input_token_num >= self.engine.context_len: + self.engine.logger.warning( + "Requested(rid=%s) token count exceeds the model's maximum context length of %s tokens. You requested a total of %s tokens: %s tokens from the input messages and %s tokens for the completion. The max_new_tokens will be truncated to %s.", + obj.rid, + self.engine.context_len, + max_new_tokens + input_token_num, + input_token_num, + max_new_tokens, + adjusted_max_new_tokens, + ) + obj.sampling_params.update({"max_new_tokens": adjusted_max_new_tokens}) + + self._maybe_wrap_json_schema_for_reasoning(obj.sampling_params) + + sampling_params = SamplingParams(**obj.sampling_params) + sampling_params.resolve_seed(obj.rid) + sampling_params.normalize(self.engine.tokenizer) + sampling_params.verify(self.engine.model_config.vocab_size) + + # Output logprobs: two request dialects, one compute path. vLLM uses + # sampling_params.logprobs; SGLang uses GenerateReqInput.return_logprob + # (+ top_logprobs_num / logprob_start_len / token_ids_logprob). Either way + # the scheduler computes only the sampled token's logprob; the response + # dialect is chosen at render time. Gate unsupported CAPABILITIES loudly + # here rather than silently clamping the request shape. + sglang_req = bool(getattr(obj, "return_logprob", False)) + return_logprob = sampling_params.logprobs is not None or sglang_req + # Output logprobs are gated by the static server arg enable_output_logprobs + # (the sampler only gathers them when on). Reject loudly instead of + # silently returning empty logprobs when the server cannot honor it. + if return_logprob and not self.engine.server_args.enable_output_logprobs: + raise ValueError( + "logprobs were requested but the server was started without " + "enable_output_logprobs; restart with enable_output_logprobs=True " + "to return output logprobs." + ) + if sglang_req: + # vLLM top-k / full-vocab are gated in SamplingParams.verify(); gate + # the SGLang capability knobs here for parity. + if getattr(obj, "top_logprobs_num", 0): + raise ValueError( + "top_logprobs_num > 0 (output top-k logprobs) is not supported " + "yet; use top_logprobs_num=0 (the sampled token's logprob)." + ) + if (getattr(obj, "logprob_start_len", -1) or -1) >= 0: + raise ValueError( + "logprob_start_len >= 0 (prompt logprobs) is not supported yet." + ) + if getattr(obj, "token_ids_logprob", None): + raise ValueError("token_ids_logprob is not supported yet.") + logprob_start_len = -1 + top_logprobs_num = 0 + token_ids_logprob = None + + if isinstance(obj, GenerateReqInput): + return TokenizedGenerateReqInput( + obj.rid, + input_text, + input_ids, + sampling_params, + return_logprob, + logprob_start_len, + top_logprobs_num, + token_ids_logprob, + obj.stream, + bootstrap_host=obj.bootstrap_host, + bootstrap_port=obj.bootstrap_port, + bootstrap_room=obj.bootstrap_room, + input_embeds=input_embeds, + session_params=session_params, + custom_logit_processor=obj.custom_logit_processor, + return_hidden_states=obj.return_hidden_states, + created_time=time.time(), + input_multi_ids=obj.input_multi_ids, + input_extra_infos=obj.input_extra_infos, + input_ids_unpadded=input_ids_unpadded, + multimodal_inputs=multimodal_inputs, + ) + + return TokenizedEmbeddingReqInput( + obj.rid, + input_text, + input_ids, + sampling_params, + created_time=time.time(), + ) diff --git a/python/tokenspeed/runtime/engine/io_struct.py b/python/tokenspeed/runtime/engine/io_struct.py new file mode 100755 index 0000000..a6bfc0a --- /dev/null +++ b/python/tokenspeed/runtime/engine/io_struct.py @@ -0,0 +1,954 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +""" +The definition of objects transferred between different +processes (TokenizerManager, DetokenizerManager, Controller). +""" + +import copy +import uuid +from abc import ABC +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, Literal + +from tokenspeed.runtime.engine.request_types import BaseFinishReason +from tokenspeed.runtime.sampling.sampling_params import SamplingParams + + +def _require(condition: bool, message: str) -> None: + if not condition: + raise ValueError(message) + + +@dataclass +class BaseReq(ABC): + rid: str | list[str] | None = field(default=None) + http_worker_ipc: str | None = field(default=None) + + def regenerate_rid(self): + """Generate a new request ID and return it.""" + if isinstance(self.rid, list): + self.rid = [uuid.uuid4().hex for _ in range(len(self.rid))] + else: + self.rid = uuid.uuid4().hex + return self.rid + + +@dataclass +class SessionParams: + id: str | None = None + rid: str | None = None + offset: int | None = None + replace: bool | None = None + + +@dataclass +class GenerateReqInput: + # The input prompt. It can be a single prompt or a batch of prompts. + text: list[str] | str | None = None + # The token ids for text; one can specify either text or input_ids + input_ids: list[list[int]] | list[int] | None = None + input_multi_ids: list[list[int]] | list[list[int]] | None = None + # The embeddings for input_ids; one can specify either text or input_ids or input_embeds. + input_embeds: list[list[list[float]]] | list[list[float]] | None = None + # Pre-built MultimodalInputs (already produced by an upstream preprocessor, + # e.g. SMG's Rust crates/multimodal pipeline). The engine's InputProcessor + # uses this directly (it does no in-process image preprocessing). input_ids + # must already contain expanded image placeholder tokens at the right + # offsets — the gateway is responsible for that. Typed as Any to avoid a + # circular import on MultimodalInputs. + precomputed_multimodal_inputs: Any | None = None + # The sampling_params. See descriptions below. + sampling_params: list[dict] | dict | None = None + input_extra_infos: list[dict] | dict | None = None + # Optional client label for logging; defaults to `rid`. Safe to reuse. + user_rid: list[str] | str | None = None + # Routing id; always server-assigned during normalize, never caller-settable. + rid: list[str] | str | None = field(default=None, init=False) + # --- Logprob request (two dialects, one compute path) --- + # vLLM-compatible requests use ``sampling_params["logprobs"]``; + # SGLang-compatible requests use the legacy fields below. A request uses + # one dialect; the response is rendered to match (override with + # ``logprob_format``). + return_logprob: list[bool] | bool | None = None + # Start location in the prompt for prompt logprobs. -1 (default) = output + # tokens only. + logprob_start_len: list[int] | int | None = None + # Number of top logprobs per position. + top_logprobs_num: list[int] | int | None = None + # Specific token ids to score per position. + token_ids_logprob: list[list[int]] | list[int] | None = None + # Detokenize tokens in the returned logprobs. + return_text_in_logprobs: bool = False + # Output rendering dialect: "vllm" | "sglang" | "both". None = auto (match + # the request dialect: vllm if sampling_params.logprobs is set, else sglang). + logprob_format: list[str | None] | str | None = None + # 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 + + # Session info for continual prompting + session_params: list[dict] | dict | None = None + + # Custom logit processor for advanced sampling control. Must be a serialized instance + # of `CustomLogitProcessor` in python/tokenspeed/runtime/sampling/custom_logit_processor.py + # Use the processor's `to_str()` method to generate the serialized string. + custom_logit_processor: list[str | None] | str | None = None + + # Whether to return hidden states + return_hidden_states: bool = False + + # For disaggregated inference + bootstrap_host: list[str] | str | None = None + bootstrap_port: list[int] | int | None = None + bootstrap_room: list[int] | int | None = None + + def normalize_batch_and_arguments(self): + 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." + ) + + # Derive 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 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: + _require( + isinstance(self.input_embeds, list), "input_embeds should be a list." + ) + if isinstance(self.input_embeds[0][0], float): + # list[list[float]] + self.is_single = True + self.batch_size = 1 + else: + # list[list[list[float]]] + _require( + isinstance(self.input_embeds[0][0], list), + "input_embeds should be a list of float lists.", + ) + _require( + isinstance(self.input_embeds[0][0][0], float), + "input_embeds should contain floats.", + ) + self.is_single = False + self.batch_size = len(self.input_embeds) + + # Handle parallel sampling. Pop "n" out of sampling_params so the + # downstream SamplingParams(**dict) construction doesn't see it — + # "n" is a request-level fan-out knob, not a per-sample field. + if self.sampling_params is None: + self.parallel_sample_num = 1 + 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 sp in self.sampling_params[1:]: + _require( + self.parallel_sample_num == sp.get("n", 1), + "The parallel_sample_num should be the same for all samples in sample params.", + ) + + 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_multi_ids is not None: + self.input_multi_ids = [self.input_multi_ids] + if self.input_embeds is not None: + self.input_embeds = [self.input_embeds] + + # Fill in default arguments + if self.is_single: + if self.sampling_params is None: + self.sampling_params = {} + if self.rid is None: + self.rid = uuid.uuid4().hex + if self.user_rid is None: + self.user_rid = self.rid + else: + if isinstance(self.user_rid, list): + _require( + len(self.user_rid) == 1, + "user_rid list should have length 1 for single request.", + ) + self.user_rid = self.user_rid[0] + _require(isinstance(self.user_rid, str), "user_rid should be a str.") + 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 + if isinstance(self.input_extra_infos, dict): + self.input_extra_infos = [self.input_extra_infos] + else: + if self.parallel_sample_num == 1: + num = self.batch_size + else: + # Expand parallel_sample_num + num = self.batch_size * self.parallel_sample_num + + if self.sampling_params is None: + self.sampling_params = [{} for _ in range(num)] + elif not isinstance(self.sampling_params, list): + self.sampling_params = [dict(self.sampling_params) for _ in range(num)] + + if self.rid is None: + self.rid = [uuid.uuid4().hex for _ in range(num)] + else: + _require(isinstance(self.rid, list), "The rid should be a list.") + if self.user_rid is None: + self.user_rid = list(self.rid) + elif isinstance(self.user_rid, str): + self.user_rid = [self.user_rid] * num + else: + _require( + isinstance(self.user_rid, list) and len(self.user_rid) == num, + "user_rid should be a str or a list of matching length.", + ) + + if self.return_logprob is None: + self.return_logprob = [False] * num + elif not isinstance(self.return_logprob, list): + self.return_logprob = [self.return_logprob] * num + else: + _require( + self.parallel_sample_num == 1, + "return_logprob cannot be a list when n > 1.", + ) + + if self.logprob_start_len is None: + self.logprob_start_len = [-1] * num + elif not isinstance(self.logprob_start_len, list): + self.logprob_start_len = [self.logprob_start_len] * num + else: + _require( + self.parallel_sample_num == 1, + "logprob_start_len cannot be a list when n > 1.", + ) + + if self.top_logprobs_num is None: + self.top_logprobs_num = [0] * num + elif not isinstance(self.top_logprobs_num, list): + self.top_logprobs_num = [self.top_logprobs_num] * num + else: + _require( + self.parallel_sample_num == 1, + "top_logprobs_num cannot be a list when n > 1.", + ) + + 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) + ] + else: + _require( + self.parallel_sample_num == 1, + "token_ids_logprob cannot be nested lists when n > 1.", + ) + + if self.logprob_format is None or isinstance(self.logprob_format, str): + self.logprob_format = [self.logprob_format] * num + + 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 + else: + _require( + self.parallel_sample_num == 1, + "custom_logit_processor cannot be a list when n > 1.", + ) + + 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 + else: + _require( + self.parallel_sample_num == 1, + "bootstrap_host cannot be a list when n > 1.", + ) + + 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 + else: + _require( + self.parallel_sample_num == 1, + "bootstrap_port cannot be a list when n > 1.", + ) + + if self.bootstrap_room is None: + self.bootstrap_room = [None] * num + elif not isinstance(self.bootstrap_room, list): + self.bootstrap_room = [self.bootstrap_room] * num + else: + _require( + self.parallel_sample_num == 1, + "bootstrap_room cannot be a list when n > 1.", + ) + + # Other checks + if self.session_params is not None: + _require( + isinstance(self.session_params, dict) + or isinstance(self.session_params[0], dict), + "session_params should be a dict or a list of dicts.", + ) + + def regenerate_rid(self): + self.rid = uuid.uuid4().hex + return self.rid + + def __getitem__(self, i): + sub = 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, + # precomputed_multimodal_inputs is a single prompt's MM; the SMG + # path only clears is_single via n>1 (batch_size == 1), so all n + # parallel samples correctly share it. Without this the image is + # silently dropped on the n>1 fan-out (placeholders -> text path). + precomputed_multimodal_inputs=self.precomputed_multimodal_inputs, + input_multi_ids=( + self.input_multi_ids[i] if self.input_multi_ids is not None else None + ), + input_embeds=( + self.input_embeds[i] if self.input_embeds is not None else None + ), + input_extra_infos=( + self.input_extra_infos[i] + if self.input_extra_infos is not None + else None + ), + sampling_params=self.sampling_params[i], + user_rid=self.user_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, + logprob_format=self.logprob_format[i], + stream=self.stream, + log_metrics=self.log_metrics, + custom_logit_processor=( + self.custom_logit_processor[i] + if self.custom_logit_processor is not None + else None + ), + return_hidden_states=self.return_hidden_states, + # 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 + ), + ) + sub.rid = self.rid[i] + return sub + + +@dataclass +class TokenizedGenerateReqInput: + # The request id + rid: str + # The input text + input_text: str + # The input token ids + input_ids: list[int] + # The sampling parameters + sampling_params: SamplingParams + # Whether to return the sampled token's logprob for this request. + return_logprob: bool + # Internal carry-over fields kept for pipeline/PD compatibility. The vLLM + # output-logprob API only drives ``return_logprob``; InputProcessor sets + # these to neutral values (logprob_start_len=-1, top_logprobs_num=0, + # token_ids_logprob=None) since prompt logprobs, output top-k, and token-id + # logprobs are not supported. + logprob_start_len: int + top_logprobs_num: int + token_ids_logprob: list[int] + # Whether to stream output + stream: bool + + # The input embeds + input_embeds: list[list[list[float]]] | list[list[float]] | None = None + + # Session info for continual prompting + session_params: SessionParams | None = None + + # Custom logit processor for advanced sampling control. Must be a serialized instance + # of `CustomLogitProcessor` in python/tokenspeed/runtime/sampling/custom_logit_processor.py + # Use the processor's `to_str()` method to generate the serialized string. + custom_logit_processor: str | None = None + + # Whether to return hidden states + return_hidden_states: bool = False + + # Time at object instantiated + created_time: float = 0.0 + + # For disaggregated inference + bootstrap_host: str | None = None + bootstrap_port: int | None = None + bootstrap_room: int | None = None + + input_multi_ids: list[list[int]] = None + input_extra_infos: list[dict] | None = None + # Original prompt ids before multimodal pad/hash replacement. The scheduler + # uses input_ids, while detokenization must use these tokenizer-valid ids. + input_ids_unpadded: list[int] | None = None + multimodal_inputs: Any | None = None + + +@dataclass +class EmbeddingReqInput: + # The input prompt. It can be a single prompt or a batch of prompts. + text: list[str] | str | None = None + # The token ids for text; one can either specify text or input_ids. + input_ids: list[list[int]] | list[int] | None = None + # Optional client label for logging; defaults to `rid`. Safe to reuse. + user_rid: list[str] | str | None = None + # Routing id; always server-assigned during normalize, never caller-settable. + rid: list[str] | str | None = field(default=None, init=False) + # Optional placeholder so non-generation callers can still instantiate the + # shared request shape without real sampling params. + sampling_params: list[dict] | dict = None + # Optional placeholder for callers that do not provide input embeddings. + input_embeds: list[list[list[float]]] | list[list[float]] | None = None + # Whether to log metrics for this request (e.g. health_generate calls do not log metrics) + log_metrics: bool = True + + def normalize_batch_and_arguments(self): + if (self.text is None and self.input_ids is None) or ( + self.text is not None and self.input_ids is not None + ): + raise ValueError("Either text or input_ids should be provided.") + + # Derive 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) + else: + 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) + + # Fill in default arguments + if self.is_single: + if self.rid is None: + self.rid = uuid.uuid4().hex + if self.user_rid is None: + self.user_rid = self.rid + else: + if isinstance(self.user_rid, list): + _require( + len(self.user_rid) == 1, + "user_rid list should have length 1 for single request.", + ) + self.user_rid = self.user_rid[0] + _require(isinstance(self.user_rid, str), "user_rid should be a str.") + 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: + _require(isinstance(self.rid, list), "The rid should be a list.") + if self.user_rid is None: + self.user_rid = list(self.rid) + elif isinstance(self.user_rid, str): + self.user_rid = [self.user_rid] * self.batch_size + else: + _require( + isinstance(self.user_rid, list) + and len(self.user_rid) == self.batch_size, + "user_rid should be a str or a list of matching length.", + ) + + if self.sampling_params is None: + self.sampling_params = [{} for _ in range(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 __getitem__(self, i): + sub = 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, + sampling_params=self.sampling_params[i], + user_rid=self.user_rid[i], + ) + sub.rid = self.rid[i] + return sub + + +@dataclass +class TokenizedEmbeddingReqInput: + # The request id + rid: str + # The input text + input_text: str + # The input token ids + input_ids: list[int] + # Placeholder sampling params field so request metadata can share one shape + # with generation-oriented code paths. + sampling_params: SamplingParams + # Time at object instantiated + created_time: float + + +@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[list[int]] + read_offsets: list[int] + # Only used when `--skip-tokenizer-init` is on + output_ids: list[int] | None + output_multi_ids: list[int] | None + # 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]] + batch_accept_draft_tokens: list[float] + + # Store some custom information, such as decoding status in multimodal scenarios, etc. + output_extra_infos: list[dict[str, Any]] + + generated_time: int + + +@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: list[int] | None + + # 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]] + batch_accept_draft_tokens: list[float] + + # Store some custom information, such as decoding status in multimodal scenarios, etc. + output_extra_infos: list[dict[str, Any]] + + generated_time: int + + +@dataclass +class BatchEmbeddingOut: + # The request id + rids: list[str] + # The finish reason + finished_reasons: list[BaseFinishReason] + # The output embedding + embeddings: list[list[float]] | list[dict] + # Token counts + prompt_tokens: list[int] + + +@dataclass +class FlushCacheReqInput: + pass + + +@dataclass +class FlushCacheReqOutput: + success: bool + + +# How a pause should treat in-flight requests. +# - "abort": kill in-flight requests immediately, then stop admitting new ones. +# - "wait": stop admitting new ones, keep stepping until running requests drain. +# - "keep": freeze everything in place; resume picks up where it left off. +PauseMode = Literal["abort", "wait", "keep"] + + +@dataclass +class PauseSchedulerReqInput: + # See PauseMode for how each mode treats in-flight requests. + mode: PauseMode = "abort" + + +@dataclass +class PauseSchedulerReqOutput: + success: bool + message: str = "" + + +@dataclass +class ResumeSchedulerReqInput: + pass + + +@dataclass +class ResumeSchedulerReqOutput: + success: bool + message: str = "" + + +@dataclass +class IsSchedulerPausedReqInput: + pass + + +@dataclass +class IsSchedulerPausedReqOutput: + is_paused: bool + + +@dataclass +class UpdateWeightFromDiskReqInput: + # The model path with the new weights + model_path: str + # The format to load the weights + load_format: str | None = None + + +@dataclass +class UpdateWeightFromDiskReqOutput: + success: bool + message: str + # Number of paused requests during weight sync. + num_paused_requests: int | None = 0 + + +@dataclass +class UpdateWeightsFromDistributedReqInput: + name: str + dtype: str + shape: list[int] + + +@dataclass +class UpdateWeightsFromDistributedReqOutput: + success: bool + message: str + + +@dataclass +class UpdateWeightsFromTensorReqInput: + serialized_named_tensors: bytes # indeed Dict[str, torch.Tensor] + load_format: str | None + flush_cache: bool + + +@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 GetWeightsByNameReqInput: + name: str + truncate_size: int = 100 + + +@dataclass +class GetWeightsByNameReqOutput: + parameter: list + + +@dataclass +class ReleaseMemoryOccupationReqInput: + # Memory regions to release. None ⇒ all ("weights" and "kv_cache"). + tags: list[str] | None = None + + +@dataclass +class ReleaseMemoryOccupationReqOutput: + success: bool = True + message: str = "" + + +@dataclass +class ResumeMemoryOccupationReqInput: + # Memory regions to resume. None ⇒ all previously released tags. + tags: list[str] | None = None + + +@dataclass +class ResumeMemoryOccupationReqOutput: + success: bool = True + message: str = "" + + +@dataclass +class IsSleepingReqInput: + pass + + +@dataclass +class IsSleepingReqOutput: + is_sleeping: bool + + +@dataclass +class AbortReq: + # The request id + rid: str + + +@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] + + +class ExpertDistributionReq(Enum): + START_RECORD = 1 + STOP_RECORD = 2 + DUMP_RECORD = 3 + + +@dataclass +class ExpertDistributionReqOutput: + pass + + +class ProfileReqType(Enum): + START_PROFILE = 1 + STOP_PROFILE = 2 + + +@dataclass +class ProfileReq: + type: ProfileReqType + output_dir: str | None = None + start_step: int | None = None + num_steps: int | None = None + activities: list[str] | None = None + profile_by_stage: bool = False + with_stack: bool | None = None + record_shapes: bool | None = None + profile_id: str | None = None + + +@dataclass +class ProfileReqOutput: + success: bool + message: str + + +@dataclass +class ConfigureLoggingReq: + log_requests: bool | None = None + log_requests_level: int | None = None + dump_requests_folder: str | None = None + dump_requests_threshold: int | None = None + + +@dataclass +class OpenSessionReqInput: + capacity_of_str_len: int + session_id: str | None = None + + +@dataclass +class CloseSessionReqInput: + session_id: str + + +@dataclass +class OpenSessionReqOutput: + session_id: str | None + success: bool + + +@dataclass +class HealthCheckOutput: + pass + + +@dataclass +class RpcReqInput: + method: str + parameters: dict | None = None + + +@dataclass +class RpcReqOutput: + success: bool + message: str + + +@dataclass +class GetLoadReqInput(BaseReq): + pass + + +@dataclass +class GetLoadReqOutput(BaseReq): + dp_rank: int = 0 + num_reqs: int = 0 + num_waiting_reqs: int = 0 + num_pages: int = 0 + + +@dataclass +class WatchLoadUpdateReq(BaseReq): + loads: list[GetLoadReqOutput] = field(default_factory=list) + + +class BlockReqType(Enum): + BLOCK = 1 + UNBLOCK = 2 + + +@dataclass +class BlockReqInput(BaseReq): + type: BlockReqType = field(default_factory=BlockReqType.BLOCK) diff --git a/python/tokenspeed/runtime/engine/llm.py b/python/tokenspeed/runtime/engine/llm.py new file mode 100644 index 0000000..c090f33 --- /dev/null +++ b/python/tokenspeed/runtime/engine/llm.py @@ -0,0 +1,152 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Sync facade over ``AsyncLLM`` for blocking callers. + +A dedicated daemon thread owns a private ``asyncio`` event loop that +drives ``AsyncLLM.generate_request`` coroutines. Sync callers submit +work via ``asyncio.run_coroutine_threadsafe``; streaming callers +consume results through a ``queue.Queue`` with three-valued +semantics (normal item / terminal sentinel / exception) so +exceptions propagate to the caller thread instead of getting +swallowed by the driver loop. + +Bridging at this layer — rather than maintaining a separate sync +IPC client — keeps ``AsyncLLM`` single-producer and the scheduler +IPC surface single-caller, while still letting ``Engine`` expose a +blocking API to callers that cannot own an event loop themselves. +""" + +import asyncio +import queue +import threading +from collections.abc import Iterator +from typing import Any + +from tokenspeed.runtime.engine.async_llm import AsyncLLM +from tokenspeed.runtime.engine.io_struct import GenerateReqInput + +# Sentinel marking the end of a streaming bridge. Identity-compared, +# so a dict payload happening to equal this value is impossible. +_STREAM_END = object() + + +class LLM: + """Sync adapter around an existing ``AsyncLLM`` instance. + + Owns a daemon thread running its own ``asyncio`` event loop; all + coroutines are driven on that loop via + ``asyncio.run_coroutine_threadsafe``. The adapter is cheap to + construct and safe to share across caller threads — the bg loop + is single-threaded, so ordering of queued work matches the + submission order. + """ + + def __init__(self, async_llm: AsyncLLM) -> None: + self.async_llm = async_llm + self._loop = asyncio.new_event_loop() + self._thread = threading.Thread( + target=self._run_loop, + name="tokenspeed-llm-sync", + daemon=True, + ) + self._thread.start() + + def _run_loop(self) -> None: + asyncio.set_event_loop(self._loop) + self._loop.run_forever() + + def run(self, coro) -> Any: + """Submit ``coro`` to the bg loop and block until it completes. + + This is the single sync entry point used by ``Engine`` for every + blocking coroutine call (``flush_cache``, ``start_profile``, + weight sync, internal-state queries, …). It replaces the + deprecated ``asyncio.get_event_loop()`` + + ``run_until_complete(coro)`` idiom that previously lived inline + at every call site. + """ + return asyncio.run_coroutine_threadsafe(coro, self._loop).result() + + def generate(self, obj: GenerateReqInput) -> dict: + """Run ``AsyncLLM.generate_request`` to completion and return the + final dict. Used by non-streaming sync paths + (``Engine.generate`` with ``stream=False``). + """ + + async def _one() -> dict: + gen = self.async_llm.generate_request(obj) + return await gen.__anext__() + + return self.run(_one()) + + def generate_stream(self, obj: GenerateReqInput) -> Iterator[dict]: + """Run ``AsyncLLM.generate_request`` and yield each output dict + as it lands. Implements the three-valued queue bridge: + + * dicts are produced by the async generator → ``q.put(item)`` + * terminal → ``q.put(_STREAM_END)`` + * exception → ``q.put(exc)`` then ``q.put(_STREAM_END)`` + + The drain coroutine is fire-and-forget on the bg loop. If the + caller abandons the iterator early (e.g. ``break`` out of the + ``for`` loop) the drain will keep running to completion and + push items into a queue nobody reads — that is acceptable for + finite ``generate_request`` streams, which is the only kind + ``AsyncLLM`` emits today. + """ + q: queue.Queue[Any] = queue.Queue() + + async def _drain() -> None: + pending_exc: BaseException | None = None + try: + async for item in self.async_llm.generate_request(obj): + q.put(item) + except BaseException as exc: # noqa: BLE001 — propagate anything + pending_exc = exc + finally: + if pending_exc is not None: + q.put(pending_exc) + q.put(_STREAM_END) + + asyncio.run_coroutine_threadsafe(_drain(), self._loop) + + while True: + item = q.get() + if item is _STREAM_END: + return + if isinstance(item, BaseException): + raise item + yield item + + def shutdown(self) -> None: + """Stop the bg event loop and join the thread. + + Idempotent. Safe to call before ``AsyncLLM`` teardown. The + daemon flag means an accidentally-skipped shutdown will not + hang interpreter exit, but leaving the loop running past + ``AsyncLLM`` teardown will surface as "Event loop is closed" + errors on later submissions, so call this before tearing + down the engine. + """ + if not self._loop.is_closed(): + self._loop.call_soon_threadsafe(self._loop.stop) + if self._thread.is_alive(): + self._thread.join(timeout=5.0) diff --git a/python/tokenspeed/runtime/engine/logprobs.py b/python/tokenspeed/runtime/engine/logprobs.py new file mode 100644 index 0000000..30df145 --- /dev/null +++ b/python/tokenspeed/runtime/engine/logprobs.py @@ -0,0 +1,243 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Logprob assembly for the async frontend — two dialects, one compute path. + +The scheduler/sampler emit format-neutral wire arrays on ``recv_obj`` +(``recv_obj.{input,output}_{token,top}_logprobs_{val,idx}`` etc.). This +processor renders them into the ``logprobs_info`` payload the per-request +``RequestOutputCollector`` merges, in whichever dialect the request asked for: + +- ``"vllm"`` -> ``meta_info["logprobs"]`` as ``list[dict[token_id, Logprob]]`` + (one dict per generated token) plus a running ``cumulative_logprob``. +- ``"sglang"`` -> ``meta_info["output_token_logprobs"]`` (and the top-k / + token-id variants) as lists of ``(logprob, token_id, text|None)`` tuples. +- ``"both"`` -> emit both (opt-in; doubles the payload). + +Only the renderers differ; the underlying wire arrays are computed once. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from tokenspeed.runtime.engine.async_llm import AsyncLLM + from tokenspeed.runtime.engine.io_struct import BatchStrOut + + +@dataclass +class Logprob: + """Per-output-token logprob entry (vLLM-style). + + Attributes: + logprob: log-probability of the sampled token. + rank: slot rank of the entry; 0 for the sampled token. NOTE: this is the + slot index, not the token's rank in the full-vocab distribution. + """ + + logprob: float + rank: int = 0 + + def to_dict(self) -> dict: + """JSON-safe view for serving boundaries that can't ship dataclasses.""" + return {"logprob": self.logprob, "rank": self.rank} + + +class LogprobsProcessor: + """Render sampler logprob wire arrays into per-request meta_info entries. + + Holds an engine reference for the live ``tokenizer`` used when the SGLang + dialect requests text decoding (``return_text=True``). The vLLM dialect does + not detokenize, so the tokenizer is never touched there. + """ + + def __init__(self, engine: AsyncLLM) -> None: + self.engine = engine + + def convert_logprob_style( + self, + logprobs_info: dict, + fmt: str, + top_logprobs_num: int, + token_ids_logprob: list[int] | None, + return_text: bool, + recv_obj: BatchStrOut, + recv_obj_index: int, + ) -> None: + """Render ``recv_obj``'s logprob arrays into ``logprobs_info``. + + ``fmt`` selects the dialect: ``"vllm"``, ``"sglang"``, or ``"both"``. + Lists EXTEND across streamed frames, so this may be called repeatedly + for one request. + """ + if fmt in ("vllm", "both"): + self._render_vllm(logprobs_info, recv_obj, recv_obj_index) + if fmt in ("sglang", "both"): + self._render_sglang( + logprobs_info, + top_logprobs_num, + token_ids_logprob, + return_text, + recv_obj, + recv_obj_index, + ) + + # --- shared wire access ------------------------------------------------- + + @staticmethod + def _row(recv_obj, field: str, idx: int): + # Defensive: sampler may not have populated logprobs for this request + # (e.g. backend doesn't support logprobs, overlap race). Treat missing + # or out-of-range wire fields as empty rather than crashing the loop. + lst = getattr(recv_obj, field, None) or [] + return lst[idx] if idx < len(lst) else [] + + # --- vLLM renderer ------------------------------------------------------ + + def _render_vllm( + self, logprobs_info: dict, recv_obj: BatchStrOut, idx: int + ) -> None: + """Emit ``logprobs`` (list[dict[int, Logprob]]) + ``cumulative_logprob``. + + Only the sampled token's logprob is materialized (rank 0). + """ + out_sampled_val = self._row(recv_obj, "output_token_logprobs_val", idx) + out_sampled_idx = self._row(recv_obj, "output_token_logprobs_idx", idx) + positions = [ + {int(out_sampled_idx[p]): Logprob(logprob=float(out_sampled_val[p]))} + for p in range(len(out_sampled_idx)) + ] + logprobs_info.setdefault("logprobs", []).extend(positions) + logprobs_info["cumulative_logprob"] = logprobs_info.get( + "cumulative_logprob", 0.0 + ) + (float(sum(out_sampled_val)) if out_sampled_val else 0.0) + + # --- SGLang renderer ---------------------------------------------------- + + def _render_sglang( + self, + logprobs_info: dict, + top_logprobs_num: int, + token_ids_logprob: list[int] | None, + return_text: bool, + recv_obj: BatchStrOut, + idx: int, + ) -> None: + """Emit the SGLang tuple-list keys (``{input,output}_token_logprobs``, + and the top-k / token-id variants when requested).""" + + def _get(field: str): + return self._row(recv_obj, field, idx) + + input_token_logprobs = logprobs_info.get("input_token_logprobs", []) + output_token_logprobs = logprobs_info.get("output_token_logprobs", []) + input_token_logprobs.extend( + self.detokenize_logprob_tokens( + _get("input_token_logprobs_val"), + _get("input_token_logprobs_idx"), + return_text, + ) + ) + output_token_logprobs.extend( + self.detokenize_logprob_tokens( + _get("output_token_logprobs_val"), + _get("output_token_logprobs_idx"), + return_text, + ) + ) + logprobs_info["input_token_logprobs"] = input_token_logprobs + logprobs_info["output_token_logprobs"] = output_token_logprobs + + if top_logprobs_num > 0: + input_top_logprobs = logprobs_info.get("input_top_logprobs", []) + output_top_logprobs = logprobs_info.get("output_top_logprobs", []) + input_top_logprobs.extend( + self.detokenize_top_logprobs_tokens( + _get("input_top_logprobs_val"), + _get("input_top_logprobs_idx"), + return_text, + ) + ) + output_top_logprobs.extend( + self.detokenize_top_logprobs_tokens( + _get("output_top_logprobs_val"), + _get("output_top_logprobs_idx"), + return_text, + ) + ) + logprobs_info["input_top_logprobs"] = input_top_logprobs + logprobs_info["output_top_logprobs"] = output_top_logprobs + + if token_ids_logprob is not None: + input_token_ids_logprobs = logprobs_info.get("input_token_ids_logprobs", []) + output_token_ids_logprobs = logprobs_info.get( + "output_token_ids_logprobs", [] + ) + input_token_ids_logprobs.extend( + self.detokenize_top_logprobs_tokens( + _get("input_token_ids_logprobs_val"), + _get("input_token_ids_logprobs_idx"), + return_text, + ) + ) + output_token_ids_logprobs.extend( + self.detokenize_top_logprobs_tokens( + _get("output_token_ids_logprobs_val"), + _get("output_token_ids_logprobs_idx"), + return_text, + ) + ) + logprobs_info["input_token_ids_logprobs"] = input_token_ids_logprobs + logprobs_info["output_token_ids_logprobs"] = output_token_ids_logprobs + + def detokenize_logprob_tokens( + self, + token_logprobs_val: list[float], + token_logprobs_idx: list[int], + decode_to_text: bool, + ): + if not decode_to_text: + return [ + (logprob, token_id, None) + for logprob, token_id in zip(token_logprobs_val, token_logprobs_idx) + ] + if self.engine.tokenizer is None: + raise RuntimeError("Tokenizer is required to decode logprob tokens.") + token_texts = self.engine.tokenizer.batch_decode(token_logprobs_idx) + return list(zip(token_logprobs_val, token_logprobs_idx, token_texts)) + + def detokenize_top_logprobs_tokens( + self, + token_logprobs_val: list, + token_logprobs_idx: list, + decode_to_text: bool, + ): + # One [k] entry per position (batch all top-k tokens across positions). + ret = [] + for logprobs, token_ids in zip(token_logprobs_val, token_logprobs_idx): + if logprobs: + ret.append( + self.detokenize_logprob_tokens(logprobs, token_ids, decode_to_text) + ) + else: + ret.append(None) + return ret diff --git a/python/tokenspeed/runtime/engine/memory_occupation.py b/python/tokenspeed/runtime/engine/memory_occupation.py new file mode 100644 index 0000000..af74455 --- /dev/null +++ b/python/tokenspeed/runtime/engine/memory_occupation.py @@ -0,0 +1,178 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""GPU-memory data plane for sleep/wake (release/resume_memory_occupation). + +Pairs with the control-plane :class:`PauseController`. A *release* pauses and +drains the scheduler (delegated to the pause controller) and only then frees GPU +memory via the torch_memory_saver adapter; a *resume* re-maps memory and, when +the KV region returns, repairs the KV cache. Tags (``weights``, ``kv_cache``) +are freed/restored independently so the RL loop can resume weights, push fresh +weights, then resume the KV cache. + +This module knows nothing about ZMQ beyond a ``send_func`` reply socket, and +nothing about scheduling beyond delegating drain to the pause controller. +""" + +from __future__ import annotations + +from collections.abc import Callable + +from tokenspeed.runtime.engine.io_struct import ( + IsSleepingReqInput, + IsSleepingReqOutput, + ReleaseMemoryOccupationReqInput, + ReleaseMemoryOccupationReqOutput, + ResumeMemoryOccupationReqInput, + ResumeMemoryOccupationReqOutput, +) +from tokenspeed.runtime.engine.pause import PauseController + +VALID_TAGS = ("weights", "kv_cache") + + +def _normalize_tags(tags: list[str] | None) -> tuple[list[str] | None, str | None]: + """Return ``(ordered_tags, error)``. ``None`` ⇒ all tags. The order is fixed + (``weights`` before ``kv_cache``) for determinism; an unknown tag yields an + error and ``None`` tags.""" + if tags is None: + return list(VALID_TAGS), None + bad = [t for t in tags if t not in VALID_TAGS] + if bad: + return None, f"invalid tags: {bad!r}; valid: {list(VALID_TAGS)}" + return [t for t in VALID_TAGS if t in tags], None + + +class MemoryOccupationController: + """Owns the data-plane release/resume for one scheduler event loop.""" + + def __init__( + self, + *, + send_func, + pause_controller: PauseController, + adapter, + enabled: bool, + reset_caches_fn: Callable[[], None], + kv_repair_fn: Callable[[], None], + kv_cache_release_allowed: bool = True, + ) -> None: + self._send = send_func + self._pause = pause_controller + self._adapter = adapter + self._enabled = enabled + self._reset_caches = reset_caches_fn + self._kv_repair = kv_repair_fn + # Whether releasing the ``kv_cache`` region is safe in this engine. False + # when prefix caching is on but the scheduler exposes no prefix-cache + # reset: discarded KV would leave stale cache entries that survive wake + # and produce silently wrong hits, so we reject the release up front + # rather than free KV we cannot invalidate. A declared capability (set + # at construction), not a runtime duck-typed probe. + self._kv_cache_release_allowed = kv_cache_release_allowed + self.released_tags: set[str] = set() + + @property + def is_sleeping(self) -> bool: + """True while any GPU-memory region is released (data-plane sleep).""" + return bool(self.released_tags) + + # -- control-request handlers (driven by the request handler) ------------- + + def handle_release(self, req: ReleaseMemoryOccupationReqInput) -> None: + if not self._enabled: + self._fail_release("memory saver not enabled (--enable-memory-saver)") + return + tags, err = _normalize_tags(req.tags) + if err is not None: + self._fail_release(err) + return + if "kv_cache" in tags and not self._kv_cache_release_allowed: + self._fail_release( + "cannot release kv_cache: prefix caching is enabled but the " + "scheduler has no prefix-cache reset, so discarded KV would " + "leave stale cache entries after wake" + ) + return + already = [t for t in tags if t in self.released_tags] + if already: + self._fail_release(f"tags already released: {already!r}") + return + if self._pause.is_drain_pending: + self._fail_release("a pause or release is already in progress") + return + # Defer the actual free until the scheduler drains. wait-mode: let + # in-flight requests finish (the RL rollout is idle between steps). + self._pause.request_drain( + abort_inflight=False, + on_drained=lambda: self._finish_release(tags), + on_cancelled=lambda: self._fail_release("resumed before release drained"), + ) + + def _finish_release(self, tags: list[str]) -> None: + if "kv_cache" in tags: + # KV is discarded; any prefix-cache entry pointing at it is stale. + self._reset_caches() + for tag in tags: + self._adapter.pause(tag=tag) + self.released_tags.add(tag) + self._pause.set_released(True) + self._send.send_pyobj(ReleaseMemoryOccupationReqOutput(success=True)) + + def _fail_release(self, message: str) -> None: + self._send.send_pyobj( + ReleaseMemoryOccupationReqOutput(success=False, message=message) + ) + + def handle_resume(self, req: ResumeMemoryOccupationReqInput) -> None: + # ``None`` on resume means "wake exactly what is asleep" — NOT all valid + # tags. Expanding to all (as on release) would make a bare resume after a + # partial release (e.g. weights-only) fail on the never-released tag, + # stranding the released region. Resolve against ``released_tags`` so a + # partial release round-trips with a default resume_memory_occupation(). + if req.tags is None: + tags = [t for t in VALID_TAGS if t in self.released_tags] + else: + tags, err = _normalize_tags(req.tags) + if err is not None: + self._send.send_pyobj( + ResumeMemoryOccupationReqOutput(success=False, message=err) + ) + return + not_released = [t for t in tags if t not in self.released_tags] + if not_released: + self._send.send_pyobj( + ResumeMemoryOccupationReqOutput( + success=False, message=f"tags not released: {not_released!r}" + ) + ) + return + for tag in tags: + self._adapter.resume(tag=tag) + self.released_tags.discard(tag) + if "kv_cache" in tags: + self._kv_repair() + if not self.released_tags: + # Fully awake → resume scheduling. + self._pause.set_released(False) + self._send.send_pyobj(ResumeMemoryOccupationReqOutput(success=True)) + + def handle_is_sleeping(self, req: IsSleepingReqInput) -> None: + self._send.send_pyobj(IsSleepingReqOutput(is_sleeping=self.is_sleeping)) diff --git a/python/tokenspeed/runtime/engine/output_processor.py b/python/tokenspeed/runtime/engine/output_processor.py new file mode 100644 index 0000000..1fbc9b9 --- /dev/null +++ b/python/tokenspeed/runtime/engine/output_processor.py @@ -0,0 +1,411 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Per-request output state and batch-output handling for the async frontend. + +Hosts: + +* ``ReqState`` — per-request bookkeeping that ``AsyncLLM`` keeps in + its ``rid_to_state`` map. +* ``OutputProcessor`` — owns the hot-path translation from scheduler + output frames (``BatchStrOut`` / ``BatchTokenIDOut`` / + ``BatchEmbeddingOut``) into the dict- + shaped payload the per-request ``RequestOutputCollector`` merges. + Also owns logprob detokenization, per-request streaming metrics, + and request dumping. Stop authority stays with the scheduler — + finish reasons are consumed as input flags, not invented here. +""" + +from __future__ import annotations + +import asyncio +import dataclasses +import logging +import time +from datetime import datetime +from pathlib import Path +from typing import TYPE_CHECKING, Any + +from tokenspeed.runtime.engine.collector import RequestOutputCollector +from tokenspeed.runtime.engine.detokenizer import IncrementalDetokenizer +from tokenspeed.runtime.engine.io_struct import ( + BatchEmbeddingOut, + BatchStrOut, + BatchTokenIDOut, +) +from tokenspeed.runtime.engine.logprobs import LogprobsProcessor +from tokenspeed.runtime.metrics.collector import RequestFinishStats + +if TYPE_CHECKING: + from tokenspeed.runtime.engine.async_llm import AsyncLLM + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class ReqState: + """Store the state a request.""" + + collector: RequestOutputCollector + finished: bool + event: asyncio.Event + obj: Any + + # For metrics + created_time: float + tokenized_time: float = 0.0 + finished_time: float = 0.0 + first_token_time: float = 0.0 + first_completion_tokens: int = 1 + last_time: float = 0.0 + last_pure_time: float = 0.0 + last_completion_tokens: int = 1 + + # For streaming output + last_output_offset: int = 0 + + # For incremental state update. + text: str = "" + output_ids: list[int] = dataclasses.field(default_factory=list) + logprobs_info: dict = dataclasses.field(default_factory=dict) + + # Inline detokenizer: lazily constructed on the first + # BatchTokenIDOut frame for this request. Stays None for + # raw-token mode (skip_tokenizer_init or tokenizer absent). + # See runtime/engine/detokenizer.py::IncrementalDetokenizer. + inline_detokenizer: IncrementalDetokenizer | None = None + + +class OutputProcessor: + """Translate scheduler output frames into per-request collector payloads. + + Owns the batch-output dispatch, logprob detokenization, streaming + metrics collection, and request dumping. The engine reference lets + this class read ``rid_to_state``, ``tokenizer``, ``server_args``, + and the metrics / dump-state fields that live on ``AsyncLLM`` + without cloning them here. + """ + + def __init__(self, engine: AsyncLLM): + self.engine = engine + self.logprobs_processor = LogprobsProcessor(engine) + + def handle_batch_output( + self, + recv_obj: BatchStrOut | BatchEmbeddingOut | BatchTokenIDOut, + ): + for i, rid in enumerate(recv_obj.rids): + state: ReqState = self.engine.rid_to_state.get(rid, None) + if state is None: + logger.error( + "Received output for rid=%r but the state was deleted in AsyncLLM.", + rid, + ) + continue + + # Build meta_info and return value + meta_info = { + "id": rid, + "finish_reason": recv_obj.finished_reasons[i], + "prompt_tokens": recv_obj.prompt_tokens[i], + } + logprobs_info = state.logprobs_info if not state.obj.stream else {} + + obj = state.obj + sp = getattr(obj, "sampling_params", None) or {} + vllm_req = sp.get("logprobs") is not None + sglang_req = bool(getattr(obj, "return_logprob", False)) + if vllm_req or sglang_req: + # Render the dialect the request asked for; default = match the + # request (vLLM via sampling_params.logprobs, else SGLang). + fmt = getattr(obj, "logprob_format", None) or ( + "vllm" if vllm_req else "sglang" + ) + try: + self.logprobs_processor.convert_logprob_style( + logprobs_info, + fmt, + getattr(obj, "top_logprobs_num", 0) or 0, + getattr(obj, "token_ids_logprob", None), + bool(getattr(obj, "return_text_in_logprobs", False)), + recv_obj, + i, + ) + meta_info.update(logprobs_info) + except Exception as exc: + logger.warning( + "Failed to attach logprobs for rid=%s: %s. Returning response without logprobs.", + rid, + exc, + ) + + if not isinstance(recv_obj, BatchEmbeddingOut): + meta_info.update( + { + "completion_tokens": recv_obj.completion_tokens[i], + "cached_tokens": recv_obj.cached_tokens[i], + } + ) + + if getattr(recv_obj, "output_hidden_states", None): + meta_info["hidden_states"] = recv_obj.output_hidden_states[i] + + if isinstance(recv_obj, BatchStrOut): + if len(recv_obj.batch_accept_draft_tokens) > 0: + meta_info.update( + {"accept_draft_tokens": recv_obj.batch_accept_draft_tokens[i]} + ) + state.text += recv_obj.output_strs[i] + if state.obj.stream: + state.logprobs_info = logprobs_info + state.output_ids.extend(recv_obj.output_ids[i]) + output_token_ids = state.output_ids[state.last_output_offset :] + state.last_output_offset = len(state.output_ids) + else: + state.logprobs_info.update(logprobs_info) + state.output_ids.extend(recv_obj.output_ids[i]) + output_token_ids = state.output_ids.copy() + + out_dict = { + "text": state.text, + "output_ids": output_token_ids, + "meta_info": meta_info, + } + if len(recv_obj.output_extra_infos): + out_dict["output_extra_info"] = recv_obj.output_extra_infos[i] + elif isinstance(recv_obj, BatchTokenIDOut): + if ( + self.engine.server_args.enable_inline_detokenizer + and self.engine.tokenizer is not None + ): + # Inline detokenizer path: run + # IncrementalDetokenizer per request and produce + # a BatchStrOut-shaped out_dict that + # RequestOutputCollector merges. + if state.inline_detokenizer is None: + state.inline_detokenizer = IncrementalDetokenizer( + decoded_text=recv_obj.decoded_texts[i], + read_offset=recv_obj.read_offsets[i], + ) + incremental_emit = state.inline_detokenizer.process( + self.engine.tokenizer, + new_decode_ids=recv_obj.decode_ids[i], + finished_reason=recv_obj.finished_reasons[i], + no_stop_trim=recv_obj.no_stop_trim[i], + skip_special_tokens=recv_obj.skip_special_tokens[i], + spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[ + i + ], + ) + if len(recv_obj.batch_accept_draft_tokens) > 0: + meta_info.update( + { + "accept_draft_tokens": recv_obj.batch_accept_draft_tokens[ + i + ] + } + ) + state.text += incremental_emit + if state.obj.stream: + state.logprobs_info = logprobs_info + state.output_ids.extend(recv_obj.decode_ids[i]) + output_token_ids = state.output_ids[state.last_output_offset :] + state.last_output_offset = len(state.output_ids) + else: + state.logprobs_info.update(logprobs_info) + state.output_ids.extend(recv_obj.decode_ids[i]) + output_token_ids = state.output_ids.copy() + + out_dict = { + "text": state.text, + "output_ids": output_token_ids, + "meta_info": meta_info, + } + if len(recv_obj.output_extra_infos): + out_dict["output_extra_info"] = recv_obj.output_extra_infos[i] + else: + # Raw-token path: skip_tokenizer_init, or + # ``enable_inline_detokenizer`` is on but + # ``self.tokenizer is None`` unexpectedly. Keep the + # response shape aligned with the BatchStrOut path by + # always populating ``text`` from the accumulated state. + if ( + self.engine.server_args.enable_inline_detokenizer + and self.engine.tokenizer is None + and not self.engine.server_args.skip_tokenizer_init + ): + logger.warning( + "AsyncLLM raw-token branch fired with " + "enable_inline_detokenizer=True and " + "skip_tokenizer_init=False; " + "self.tokenizer is unexpectedly None. " + "Output text will be empty for rid=%s.", + rid, + ) + + output_multi_ids = None + if self.engine.server_args.stream_output and state.obj.stream: + state.output_ids.extend(recv_obj.output_ids[i]) + output_token_ids = state.output_ids[state.last_output_offset :] + if recv_obj.output_multi_ids is not None: + output_multi_ids = recv_obj.output_multi_ids[i][ + state.last_output_offset : + ] + state.last_output_offset = len(state.output_ids) + else: + state.output_ids.extend(recv_obj.output_ids[i]) + output_token_ids = state.output_ids.copy() + if recv_obj.output_multi_ids is not None: + output_multi_ids = recv_obj.output_multi_ids[i] + + if len(recv_obj.batch_accept_draft_tokens) > 0: + meta_info.update( + { + "accept_draft_tokens": recv_obj.batch_accept_draft_tokens[ + i + ] + } + ) + + out_dict = { + "text": state.text, + "output_ids": output_token_ids, + "meta_info": meta_info, + } + if len(recv_obj.output_extra_infos): + out_dict["output_extra_info"] = recv_obj.output_extra_infos[i] + if output_multi_ids is not None: + out_dict["output_multi_ids"] = output_multi_ids + else: + out_dict = { + "embedding": recv_obj.embeddings[i], + "meta_info": meta_info, + } + + state.finished = recv_obj.finished_reasons[i] is not None + if state.finished: + if self.engine.server_args.speculative_algorithm: + meta_info["spec_verify_ct"] = recv_obj.spec_verify_ct[i] + state.finished_time = time.time() + meta_info["e2e_latency"] = state.finished_time - state.created_time + + state.collector.put( + out_dict, stream=bool(getattr(state.obj, "stream", False)) + ) + state.event.set() + + # Log metrics and dump + if self.engine.enable_metrics and not isinstance( + recv_obj, BatchEmbeddingOut + ): + self.collect_metrics(state, recv_obj, i) + if ( + self.engine.dump_requests_folder + and state.finished + and state.obj.log_metrics + ): + self.dump_requests(state, out_dict) + + def collect_metrics(self, state: ReqState, recv_obj, i: int): + completion_tokens = ( + recv_obj.completion_tokens[i] + if getattr(recv_obj, "completion_tokens", None) + else 0 + ) + + if state.first_token_time == 0.0: + state.first_token_time = state.last_time = time.time() + state.last_pure_time = recv_obj.generated_time + state.last_completion_tokens = completion_tokens + state.first_completion_tokens = completion_tokens + self.engine.metrics.observe_time_to_first_token( + state.first_token_time - state.created_time + ) + else: + num_new_tokens = completion_tokens - state.last_completion_tokens + if num_new_tokens: + new_time = time.time() + interval = new_time - state.last_time + pure_interval = recv_obj.generated_time - state.last_pure_time + self.engine.metrics.observe_inter_token_latency( + interval, + num_new_tokens, + ) + self.engine.metrics.observe_inter_token_latency( + pure_interval, num_new_tokens + ) + state.last_pure_time = recv_obj.generated_time + state.last_time = new_time + state.last_completion_tokens = completion_tokens + + if state.finished: + fr = recv_obj.finished_reasons[i] + # TODO: consolidate the return type of fr. + finished_ok = not ( + fr.get("type") == "abort" + if isinstance(fr, dict) + else getattr(fr, "is_error", False) + ) + cached_prompt = ( + recv_obj.cached_tokens[i] + if getattr(recv_obj, "cached_tokens", None) is not None + else 0 + ) + self.engine.metrics.record_request_finish( + RequestFinishStats( + prompt_tokens=recv_obj.prompt_tokens[i], + generation_tokens=completion_tokens, + e2e_latency=state.finished_time - state.created_time, + cached_prompt_tokens=cached_prompt, + finished_ok=finished_ok, + ) + ) + if (completion_tokens - state.first_completion_tokens) > 0: + self.engine.metrics.observe_inter_token_latency( + state.finished_time - state.first_token_time, + completion_tokens - state.first_completion_tokens, + ) + + def dump_requests(self, state: ReqState, out_dict: dict): + import pickle as _pickle + + self.engine.dump_request_list.append( + (state.obj, out_dict, state.created_time, time.time()) + ) + + if len(self.engine.dump_request_list) >= self.engine.dump_requests_threshold: + dump_folder = Path(self.engine.dump_requests_folder) + filename = dump_folder / ( + datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + ".pkl" + ) + logger.info( + "Dump %s requests to %s", len(self.engine.dump_request_list), filename + ) + + to_dump = self.engine.dump_request_list + self.engine.dump_request_list = [] + + def background_task(): + dump_folder.mkdir(parents=True, exist_ok=True) + with filename.open("wb") as dump_file: + _pickle.dump(to_dump, dump_file) + + # Schedule the task to run in the background without awaiting it + asyncio.create_task(asyncio.to_thread(background_task)) diff --git a/python/tokenspeed/runtime/engine/parallel_sampling.py b/python/tokenspeed/runtime/engine/parallel_sampling.py new file mode 100644 index 0000000..8422206 --- /dev/null +++ b/python/tokenspeed/runtime/engine/parallel_sampling.py @@ -0,0 +1,97 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Tokenized-payload prep for parallel-sampling (n>1) fan-out. + +The scheduler treats each replica as a standalone rid (the n>1 expansion +happens at the AsyncLLM frontend, not on the scheduler), so there is no +per-replica scheduler state here. What we need to centralize is the +tokenized-payload mutation the fan-out loop performs on every replica: +rid regeneration, the ``max_new_tokens=0`` warmup variant, and the plain +replica copy. + +The two functions below are pure (no engine state, no I/O) and are +called from ``AsyncLLM._handle_batch_request`` once for the prefix +warmup and ``parallel_sample_num`` times for the replica expansion. +""" + +from __future__ import annotations + +import copy + +from tokenspeed.runtime.engine.io_struct import ( + EmbeddingReqInput, + GenerateReqInput, + TokenizedEmbeddingReqInput, + TokenizedGenerateReqInput, +) + + +def _own_multimodal_inputs(tokenized_copy) -> None: + """Give this fan-out copy its own multimodal feature tensors. + + ``MultimodalInputs.publish_shm_features`` rewrites ``item.feature`` + in place from tensor to SHM handle, and consumers later unlink the + segment. The prefix warmup and each n>1 replica are shallow + ``copy.copy`` of the same tokenized payload, so without an independent + copy they would publish only once (the 2nd+ ``publish`` is a no-op + because the feature is already a handle) and share the SHM segment. + When the first replica's scheduler step consumes-and-unlinks, the + later replicas hit ``FileNotFoundError`` on their own attach. + """ + mm = getattr(tokenized_copy, "multimodal_inputs", None) + if mm is not None: + tokenized_copy.multimodal_inputs = copy.deepcopy(mm) + + +def prepare_prefix_warmup( + tmp_obj: GenerateReqInput | EmbeddingReqInput, + tokenized_obj: TokenizedGenerateReqInput | TokenizedEmbeddingReqInput, +) -> TokenizedGenerateReqInput | TokenizedEmbeddingReqInput: + """Build the prefix-warmup variant used before parallel-sampling + fan-out. Mutates ``tmp_obj`` to receive a fresh rid; returns a + copy of ``tokenized_obj`` with ``max_new_tokens`` forced to 0 + and streaming disabled so the scheduler caches the common + prefix before the replicas are dispatched. + """ + tokenized_copy = copy.copy(tokenized_obj) + _own_multimodal_inputs(tokenized_copy) + tokenized_copy.rid = tmp_obj.regenerate_rid() + tokenized_copy.sampling_params = copy.copy(tokenized_copy.sampling_params) + tokenized_copy.sampling_params.max_new_tokens = 0 + tokenized_copy.stream = False + return tokenized_copy + + +def prepare_parallel_sampling_replica( + tmp_obj: GenerateReqInput | EmbeddingReqInput, + tokenized_obj: TokenizedGenerateReqInput | TokenizedEmbeddingReqInput, +) -> TokenizedGenerateReqInput | TokenizedEmbeddingReqInput: + """Build one tokenized replica for parallel-sampling fan-out. + + Mutates ``tmp_obj`` to receive a fresh rid; returns a copy of + ``tokenized_obj`` sharing that rid. The rest of the tokenized + payload (sampling_params, input_ids, etc.) is unchanged because + the replicas share everything except their request identity. + """ + tokenized_copy = copy.copy(tokenized_obj) + _own_multimodal_inputs(tokenized_copy) + tokenized_copy.rid = tmp_obj.regenerate_rid() + return tokenized_copy diff --git a/python/tokenspeed/runtime/engine/pause.py b/python/tokenspeed/runtime/engine/pause.py new file mode 100644 index 0000000..c9c5696 --- /dev/null +++ b/python/tokenspeed/runtime/engine/pause.py @@ -0,0 +1,311 @@ +# SPDX-License-Identifier: Apache-2.0 +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Pause / resume control state for a scheduler event loop. + +The pause gate lives in Python: requests are admitted to the scheduler from the +event loop (``scheduler.submit_requests``), so withholding new work while paused +is handled here rather than inside the scheduler itself. + +Modes (how a pause treats in-flight requests): + +- ``abort``: cancel in-flight requests, then drain and reply. +- ``wait`` : let in-flight requests finish naturally, then drain and reply. +- ``keep`` : freeze everything in place; reply immediately; resume later. + +``abort`` and ``wait`` both leave the scheduler in ``PAUSED_NEW`` (no new +requests admitted, running requests keep stepping) and defer their reply until +the scheduler has drained. ``keep`` moves to ``PAUSED_ALL`` (nothing scheduled) +and replies immediately. +""" + +from __future__ import annotations + +import enum +from collections.abc import Callable +from dataclasses import dataclass + +from tokenspeed.runtime.engine.io_struct import ( + IsSchedulerPausedReqInput, + IsSchedulerPausedReqOutput, + PauseSchedulerReqInput, + PauseSchedulerReqOutput, + ResumeSchedulerReqInput, + ResumeSchedulerReqOutput, +) + + +@dataclass +class _PendingDrain: + """A deferred action resolved when the scheduler drains. + + ``on_drained`` runs once ``scheduler_drained`` is true: it sends the success + reply and, for a memory release, frees GPU memory. ``on_cancelled`` runs if a + resume arrives before the drain completes: it sends the failure reply to the + correct communicator (pause vs release use different ZMQ channels, so the + action carries its own reply rather than the controller hard-coding one). + """ + + on_drained: Callable[[], None] + on_cancelled: Callable[[], None] + + +class PauseState(enum.IntEnum): + """Scheduler pause state. + + - ``UNPAUSED``: normal operation. + - ``PAUSED_NEW``: no new requests admitted; running requests keep stepping. + - ``PAUSED_ALL``: nothing scheduled; everything frozen in place. + """ + + UNPAUSED = 0 + PAUSED_NEW = 1 + PAUSED_ALL = 2 + + +def scheduler_drained(scheduler) -> bool: + """True when the scheduler holds no requests that need a forward pass. + + Covers every active lifecycle state (waiting/submitted, prefilling, + decoding, retracted). Post-finish writeback states are async teardown and + do not run forward work, so they do not block a drain. + """ + return ( + scheduler.waiting_size() == 0 + and scheduler.decoding_size() == 0 + and scheduler.prefilling_size() == 0 + and scheduler.retract_count() == 0 + ) + + +class PauseController: + """Owns pause/resume state for one scheduler event loop. + + Split of responsibilities: the ``handle_*`` methods are driven by the + request handler (they only touch local state + the reply socket); the event + loop queries :pyattr:`admit_blocked` / :pyattr:`forward_blocked`, drains + buffered specs on resume, and calls :pymeth:`maybe_finish_drain` once per + iteration to resolve a deferred abort/wait reply. + + ``send_func`` is the scheduler→tokenizer reply socket (a no-op + ``_NullSender`` on non-rank-0 TP ranks, matching the existing control-reply + pattern), so the ``handle_*`` methods are safe to call on every rank. + """ + + def __init__(self, send_func) -> None: + self._send = send_func + self.state = PauseState.UNPAUSED + # RequestSpecs withheld from the scheduler while paused; flushed on resume. + self.buffered_specs: list = [] + # Deferred post-drain action for abort/wait pause OR memory release; held + # until the scheduler drains. Single-consumer: only one may be armed. + self._pending_drain: _PendingDrain | None = None + # True once GPU memory has actually been released (data plane). Distinct + # from forward_blocked: PAUSED_ALL alone still permits DP idle forwards, + # which run the model (touch weights) and must be suppressed while the + # weights region is unmapped. + self.released: bool = False + # Set by a pause(mode="abort"); consumed once by the event loop to + # cancel in-flight requests already in the scheduler. + self._abort_all_pending = False + # Set by abort/wait; consumed once by the event loop to cancel requests + # still compiling in the grammar queue (not yet in the scheduler). + self._cancel_grammar_pending = False + + # -- state queried by the event loop -------------------------------------- + + @property + def is_paused(self) -> bool: + return self.state != PauseState.UNPAUSED + + @property + def admit_blocked(self) -> bool: + """Whether new requests should be withheld from the scheduler.""" + return self.state != PauseState.UNPAUSED + + @property + def forward_blocked(self) -> bool: + """Whether the loop should run no forward work this iteration.""" + return self.state == PauseState.PAUSED_ALL + + def consume_abort_all(self) -> bool: + """Return (once) whether the event loop should cancel all in-flight reqs.""" + if self._abort_all_pending: + self._abort_all_pending = False + return True + return False + + def consume_cancel_grammar(self) -> bool: + """Return (once) whether the event loop should cancel grammar-queued reqs. + + Set for abort/wait: requests still compiling a grammar are not yet in + the scheduler or ``rid_to_state``, so the abort sweep and the drain + check both miss them. Left compiling, they would be promoted and either + run after a weight swap (abort) or be buffered past a wait drain. They + have produced no output, so cancelling them is safe. + """ + if self._cancel_grammar_pending: + self._cancel_grammar_pending = False + return True + return False + + def buffer_specs(self, specs: list) -> None: + self.buffered_specs.extend(specs) + + def take_buffered_specs(self) -> list: + specs, self.buffered_specs = self.buffered_specs, [] + return specs + + # -- generic drain machinery (shared by pause and memory release) ---------- + + @property + def is_drain_pending(self) -> bool: + return self._pending_drain is not None + + def request_drain( + self, + *, + abort_inflight: bool, + on_drained: Callable[[], None], + on_cancelled: Callable[[], None], + ) -> bool: + """Start a wait-style drain (PAUSED_NEW, cancel grammar-queued) and arm a + post-drain action. Returns False if a drain is already pending (the + caller should send its own busy reply). ``abort_inflight=True`` also + cancels in-flight requests (abort mode); False lets them finish (wait + mode / memory release).""" + if self._pending_drain is not None: + return False + self.state = PauseState.PAUSED_NEW + self._pending_drain = _PendingDrain(on_drained, on_cancelled) + self._cancel_grammar_pending = True + if abort_inflight: + self._abort_all_pending = True + return True + + def set_released(self, released: bool) -> None: + """Mark GPU memory released (freeze fully) or restored (unpause).""" + self.released = released + self.state = PauseState.PAUSED_ALL if released else PauseState.UNPAUSED + + # -- control-request handlers (driven by the request handler) ------------- + + def handle_pause(self, req: PauseSchedulerReqInput) -> None: + if req.mode not in ("abort", "wait", "keep"): + self._send.send_pyobj( + PauseSchedulerReqOutput( + success=False, message=f"invalid pause mode: {req.mode!r}" + ) + ) + return + + # Reject any new pause while an abort/wait pause or memory release is + # still draining: the post-drain action is a single-consumer promise + # (``_pending_drain``), so a second drain would overwrite it and strand + # the first caller forever on its ZMQ await. ``keep`` never arms a drain, + # so it can't be the *first* pause here, but it must not clobber a + # draining one either. + if self._pending_drain is not None: + self._send.send_pyobj( + PauseSchedulerReqOutput( + success=False, message="a pause is already in progress" + ) + ) + return + + if req.mode == "keep": + # Freeze in place and reply now — nothing to drain. + self.state = PauseState.PAUSED_ALL + self._send.send_pyobj(PauseSchedulerReqOutput(success=True)) + return + + # abort / wait: stop admitting new requests, keep stepping so in-flight + # requests drain, and reply once the scheduler is empty. Both also + # cancel grammar-queued (still-compiling) pre-pause requests. + self.request_drain( + abort_inflight=(req.mode == "abort"), + on_drained=lambda: self._send.send_pyobj( + PauseSchedulerReqOutput(success=True) + ), + on_cancelled=lambda: self._send.send_pyobj( + PauseSchedulerReqOutput( + success=False, message="resumed before pause drained" + ) + ), + ) + + def handle_resume(self, req: ResumeSchedulerReqInput) -> None: + # Reject a scheduler-level resume while GPU memory is still released. + # ``released`` is owned by the memory controller (its ``set_released`` + # is the sole writer); clearing it here would flip the state to + # UNPAUSED without remapping the weights/KV regions, so the next admit + # or DP idle forward would touch unmapped memory. The caller must wake + # via ``resume_memory_occupation`` instead. + if self.released: + self._send.send_pyobj( + ResumeSchedulerReqOutput( + success=False, + message=( + "memory is released; call resume_memory_occupation to " + "wake before resuming the scheduler" + ), + ) + ) + return + # If a wait/abort pause is still awaiting its drain reply, it has NOT + # drained — ``maybe_finish_drain`` clears ``_pending_reply`` the instant + # it does. We must still reply (resume uses a separate communicator and + # cannot otherwise wake the pause caller, who would block forever), but + # the reply must be a failure: acking success here would tell a + # weight-swapping caller it is safe to proceed while pre-pause requests + # are still in flight under the old weights. + if self._pending_drain is not None: + action = self._pending_drain + self._pending_drain = None + action.on_cancelled() + # Buffered specs are flushed by the event loop on its next admission + # pass (state is already UNPAUSED by then). ``released`` is intentionally + # NOT touched here — see the guard above; only set_released() writes it. + self.state = PauseState.UNPAUSED + self._abort_all_pending = False + self._cancel_grammar_pending = False + self._send.send_pyobj(ResumeSchedulerReqOutput(success=True)) + + def handle_is_paused(self, req: IsSchedulerPausedReqInput) -> None: + self._send.send_pyobj(IsSchedulerPausedReqOutput(is_paused=self.is_paused)) + + # -- per-iteration drain check (driven by the event loop) ----------------- + + def maybe_finish_drain(self, scheduler) -> None: + """Resolve a deferred pause/release action once the scheduler has drained. + + The action is cleared *before* it runs so a release's ``on_drained`` can + re-arm controller state (``set_released``) without tripping the + single-consumer guard. + """ + if self._pending_drain is None: + return + if not scheduler_drained(scheduler): + return + action = self._pending_drain + self._pending_drain = None + action.on_drained() diff --git a/python/tokenspeed/runtime/engine/protocol.py b/python/tokenspeed/runtime/engine/protocol.py new file mode 100644 index 0000000..3b4241f --- /dev/null +++ b/python/tokenspeed/runtime/engine/protocol.py @@ -0,0 +1,228 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Serving-facing engine protocol. + +``EngineClient`` is the narrow surface that the OpenAI serving layer +and ``http_server.py`` are allowed to depend on, letting callers stop +typing against the concrete ``AsyncLLM`` class. +``AsyncLLM(SchedulerControlClient, EngineClient)`` inherits the protocol so the +conformance is a class-definition-time invariant rather than duck +typing at every call site. The protocol stays ``@runtime_checkable`` so +``isinstance(engine, EngineClient)`` remains a lightweight check +(exercised by ``test/runtime/test_async_llm_protocol.py``). + +What's on the protocol +---------------------- +The surface covers the calls that the serving layer actually makes: +the generate/embed/abort/attribute path plus the administrative +RPCs (weights, cache, session, profile, expert-distribution, +load-query, internal-state, logging config). + +What's intentionally off the protocol +------------------------------------- +Two categories stay concrete and are accessed via a narrower type +(or via ``isinstance`` casts in the caller): + +1. Attribute escape hatches — ``rid_to_state``, ``server_status``. + ``http_server.py`` reads them directly for liveness / health + reasons that are out of scope for the serving-facing protocol. + +2. Purely internal coordination state — ``model_update_lock``, + ``session_futures``, ``flush_cache_communicator`` and the other + ``SchedulerControlClient`` ``*_communicator`` attributes. + +If a caller needs any of the above, it must hold a concrete +``AsyncLLM`` reference, not an ``EngineClient``-typed one. This is +deliberate. +""" + +from collections.abc import AsyncGenerator +from typing import ( + Any, + Protocol, + runtime_checkable, +) + +from tokenspeed.runtime.configs.model_config import ModelConfig +from tokenspeed.runtime.engine.io_struct import ( + CloseSessionReqInput, + ConfigureLoggingReq, + EmbeddingReqInput, + FlushCacheReqOutput, + GenerateReqInput, + GetLoadReqOutput, + GetWeightsByNameReqInput, + InitWeightsUpdateGroupReqInput, + OpenSessionReqInput, + ReleaseMemoryOccupationReqInput, + ResumeMemoryOccupationReqInput, + SetInternalStateReq, + UpdateWeightFromDiskReqInput, + UpdateWeightsFromDistributedReqInput, + UpdateWeightsFromTensorReqInput, +) +from tokenspeed.runtime.utils.server_args import ServerArgs + + +@runtime_checkable +class EngineClient(Protocol): + """Serving-facing async engine surface. + + ``AsyncLLM(SchedulerControlClient, EngineClient)`` inherits + this protocol explicitly, so conformance is structurally + guaranteed at class-definition time rather than relying on + duck typing at every call site. + """ + + # ---- Configuration / identity --------------------------------- + # Mutable because ``update_weights_from_disk`` reassigns + # ``served_model_name`` and ``model_path`` on successful reloads + # (see ``_wait_for_model_update_from_disk``). Typed as their + # current runtime shape. + server_args: ServerArgs + model_config: ModelConfig + tokenizer: Any + model_path: str + served_model_name: str + is_generation: bool + is_image_gen: bool + + # ---- Liveness state ------------------------------------------- + # Monotonic epoch-second timestamp of the last message received + # from the scheduler's shared output socket. ``http_server.py`` + # reads this for health / idle-timeout logic. + last_receive_tstamp: float + gracefully_exit: bool + + # ---- Generate / embed path ------------------------------------ + + async def generate_request( + self, + obj: GenerateReqInput | EmbeddingReqInput, + ) -> AsyncGenerator[dict[str, Any], None]: ... + + def abort_request(self, rid: str) -> None: ... + + # ---- Session management -------------------------------------- + + async def open_session( + self, + obj: OpenSessionReqInput, + ) -> str | None: ... + + async def close_session( + self, + obj: CloseSessionReqInput, + ) -> None: ... + + # ---- Cache / logging config ---------------------------------- + + async def flush_cache(self) -> FlushCacheReqOutput: ... + + def configure_logging(self, obj: ConfigureLoggingReq) -> None: ... + + # ---- Pause / resume (RLHF weight-update control) -------------- + + async def pause_scheduler(self, *, mode: str = "abort") -> bool: ... + + async def resume_scheduler(self) -> bool: ... + + async def is_scheduler_paused(self) -> bool: ... + + # ---- Server lifecycle / health ------------------------------- + + def is_server_starting(self) -> bool: ... + + def mark_server_up(self) -> None: ... + + def mark_server_unhealthy(self) -> None: ... + + def drop_request_state(self, rid: str) -> None: ... + + # ---- Weight-update RPCs -------------------------------------- + + async def update_weights_from_disk( + self, + obj: UpdateWeightFromDiskReqInput, + ) -> tuple[bool, str, Any]: ... + + async def init_weights_update_group( + self, + obj: InitWeightsUpdateGroupReqInput, + ) -> tuple[bool, str]: ... + + async def update_weights_from_distributed( + self, + obj: UpdateWeightsFromDistributedReqInput, + ) -> tuple[bool, str]: ... + + async def update_weights_from_tensor( + self, + obj: UpdateWeightsFromTensorReqInput, + ) -> tuple[bool, str]: ... + + async def get_weights_by_name( + self, + obj: GetWeightsByNameReqInput, + ) -> Any: ... + + # ---- Memory occupation RPCs ---------------------------------- + + async def release_memory_occupation( + self, + obj: ReleaseMemoryOccupationReqInput, + ) -> None: ... + + async def resume_memory_occupation( + self, + obj: ResumeMemoryOccupationReqInput, + ) -> None: ... + + async def is_sleeping(self) -> bool: ... + + # ---- Profiling / expert distribution ------------------------- + + async def start_profile( + self, + output_dir: str | None = None, + start_step: int | None = None, + num_steps: int | None = None, + activities: list[str] | None = None, + with_stack: bool | None = None, + record_shapes: bool | None = None, + profile_by_stage: bool = False, + ) -> Any: ... + + async def stop_profile(self) -> Any: ... + + async def start_expert_distribution_record(self) -> None: ... + + async def stop_expert_distribution_record(self) -> None: ... + + async def dump_expert_distribution_record(self) -> None: ... + + # ---- Engine internal state ----------------------------------- + + async def get_internal_state(self) -> list[dict[Any, Any]]: ... + + async def set_internal_state(self, obj: SetInternalStateReq) -> list[bool]: ... + + async def get_load(self) -> list[GetLoadReqOutput]: ... diff --git a/python/tokenspeed/runtime/engine/request.py b/python/tokenspeed/runtime/engine/request.py new file mode 100755 index 0000000..3d9da79 --- /dev/null +++ b/python/tokenspeed/runtime/engine/request.py @@ -0,0 +1,353 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import copy +import time +from typing import Any + +import torch + +from tokenspeed.runtime.cache.req_to_token_pool import ( + ReqToTokenPoolInfo, +) +from tokenspeed.runtime.engine.request_types import ( # noqa: F401 + ABORT_CODE, + FINISH_ABORT, + FINISH_LENGTH, + FINISH_MATCHED_STR, + FINISH_MATCHED_TOKEN, + INIT_INCREMENTAL_DETOKENIZATION_OFFSET, + BaseFinishReason, +) +from tokenspeed.runtime.grammar.base_grammar_backend import BaseGrammarObject +from tokenspeed.runtime.metrics.collector import TimeStats +from tokenspeed.runtime.sampling.sampling_params import SamplingParams +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +class Req: + """The input and output status of a request.""" + + def __init__( + self, + rid: str, + origin_input_text: str, + origin_input_ids: tuple[int], + sampling_params: SamplingParams, + return_logprob: bool = False, + top_logprobs_num: int = 0, + token_ids_logprob: list[int] = None, + stream: bool = False, + origin_input_ids_unpadded: tuple[int] | None = None, + input_embeds: list[list[float]] | None = None, + input_extra_infos: list[dict] | None = None, + session_id: str | None = None, + custom_logit_processor: str | None = None, + return_hidden_states: bool = False, + eos_token_ids: set[int] | None = None, + bootstrap_host: str | None = None, + bootstrap_port: int | None = None, + bootstrap_room: int | None = None, + origin_input_multi_ids: list[list[int]] | None = None, + created_time: float | None = None, + ): + # Input and output info + self.rid = rid + self.origin_input_text = origin_input_text + self.origin_input_ids_unpadded = ( + origin_input_ids_unpadded + if origin_input_ids_unpadded + else origin_input_ids # Before image padding + ) + self.origin_input_ids = origin_input_ids + self.origin_input_multi_ids = origin_input_multi_ids + # Each decode stage's output ids + self.output_ids = [] + self.output_multi_ids = [] + # fill_ids = origin_input_ids + output_ids. Updated if chunked. + self.fill_ids = None + self.fill_multi_ids = None + self.fill_input_embeds = None + # For Eagle and chunked prefill, remove first token when chunked prefill + self.draft_fill_ids = None + self.session_id = session_id + self.input_embeds = input_embeds + self.input_extra_infos = input_extra_infos + + # Sampling info + if isinstance(sampling_params.custom_params, dict): + sampling_params = copy.copy(sampling_params) + sampling_params.custom_params = sampling_params.custom_params | { + "__req__": self + } + self.sampling_params = sampling_params + + self.custom_logit_processor = custom_logit_processor + self.return_hidden_states = return_hidden_states + + # Memory pool info + self.req_pool_idx: int | None = None + self.req_to_token_pool_info: ReqToTokenPoolInfo | None = None + # substitute for prefix_indices + self.prefix_page_ids = [] + self.prefix_len = 0 + + # Check finish + self.tokenizer = None + # Cached tokenizer-related ids to avoid repeated HF attribute lookups in check_finished(). + self._eos_token_id_cached: int | None = None + self._additional_stop_token_ids_cached: set[int] | None = None + self.finished_reason = None + # Whether this request has finished output + self.finished_output = None + # If we want to abort the request in the middle of the event loop, set this to true + # Note: We should never set finished_reason in the middle, the req will get filtered and never respond + self.to_abort = False + # This carries the error message for `.to_abort` and will be attached to the finished_reason at the end of the event loop + self.to_abort_message: str = "Unknown error" + self.stream = stream + self.eos_token_ids = eos_token_ids + + # For incremental decoding + # ----- | --------- read_ids -------| + # ----- | surr_ids | + # xxxxx | xxxxxxxxxxx | xxxxxxxxxxx | + # ----- ^ ----------- ^ ----------- ^ + # ----- 1 ----------- 2 ----------- 3 + # 1: surr_offset + # 2: read_offset + # 3: last token + self.surr_offset = None # Surrounding offset to defeat the cleanup algorithm + self.read_offset = None + self.decoded_text = "" + + # Prefix info + # The indices to kv cache for the shared prefix. + self.prefix_indices = [] + # Number of tokens to run prefill. + self.extend_input_len = 0 + # The relative logprob_start_len in an extend batch + self.extend_logprob_start_len = 0 + self.last_node = None + + # Whether or not if it is chunked. It increments whenever + # it is chunked, and decrement whenever chunked request is + # processed. + self.is_chunked = 0 + + # For retraction + self.is_retracted = False + + # Incremental streamining + self.send_token_offset: int = 0 + self.send_decode_id_offset: int = 0 + # because the decode server does not have the first output token logprobs + self.send_output_token_logprobs_offset: int = 0 + + # Logprobs (arguments) + self.return_logprob = return_logprob + # Start index to compute logprob from. + self.logprob_start_len = 0 + self.top_logprobs_num = top_logprobs_num + self.token_ids_logprob = token_ids_logprob + + # Logprobs (return values) + self.input_logprob_sent: bool = False + self.input_token_logprobs_val: list[float] | None = None + self.input_token_logprobs_idx: list[int] | None = None + self.input_top_logprobs_val: list[float] | None = None + self.input_top_logprobs_idx: list[int] | None = None + self.input_token_ids_logprobs_val: list[float] | None = None + self.input_token_ids_logprobs_idx: list[int] | None = None + # Temporary holder to store input_token_logprobs. + self.input_token_logprobs: list[tuple[int]] | None = None + self.temp_input_top_logprobs_val: list[torch.Tensor] | None = None + self.temp_input_top_logprobs_idx: list[int] | None = None + self.temp_input_token_ids_logprobs_val: list[float] | None = None + self.temp_input_token_ids_logprobs_idx: list[int] | None = None + + if return_logprob: + self.output_token_logprobs_val = [] + self.output_token_logprobs_idx = [] + self.output_top_logprobs_val = [] + self.output_top_logprobs_idx = [] + self.output_token_ids_logprobs_val = [] + self.output_token_ids_logprobs_idx = [] + else: + self.output_token_logprobs_val = self.output_token_logprobs_idx = ( + self.output_top_logprobs_val + ) = self.output_top_logprobs_idx = self.output_token_ids_logprobs_val = ( + self.output_token_ids_logprobs_idx + ) = None + self.hidden_states = [] + + # Embedding (return values) + self.embedding = None + + # Constrained decoding + self.grammar: BaseGrammarObject | None = None + + # The number of cached tokens that were already cached in the KV cache + self.cached_tokens = 0 + self.already_computed = 0 + self.last_host_node: Any = None + self.host_hit_length = 0 + + # The number of verification forward passes in the speculative decoding. + # This is used to compute the average acceptance length per request. + self.spec_verify_ct = 0 + + # Time of obj created + # Use the created_time from tokenizer if provided, otherwise use current time + if created_time is not None: + self.created_time = created_time + else: + self.created_time = time.time() + # Calculate the time from receiving the request at TokenizerManager to reaching process_input_requests in the scheduling process + self.tokenizer_to_scheduler_latency = time.time() - self.created_time + # For metrics + self.time_stats: TimeStats = TimeStats() + self.has_log_time_stats: bool = False + self.queue_time_start = None + self.queue_time_end = None + self.last_tic = time.monotonic() + self.first_latency_recorded = ( + False # Flag to track if first latency has been recorded + ) + self.prefill_waiting_recorded = False + self.first_chunk_forward_start_time = None + + self.reserve_num_tokens = 0 + # For disaggregation + self.bootstrap_host: str = bootstrap_host + self.bootstrap_port: int | None = bootstrap_port + self.bootstrap_room: int | None = bootstrap_room + + # the start index of the sent kv cache + # We want to send it chunk by chunk for chunked prefill. + # After every chunk forward, we do the following: + # kv_send(req.input_ids[req.start_send_idx:len(req.fill_ids)]) + # start_send_idx = len(req.fill_ids) + self.start_send_idx: int = 0 + + # For overlap schedule, we delay the kv transfer until `process_batch_result_disagg_prefill` rather than `process_prefill_chunk` in non-overlap + # This is because kv is not ready in `process_prefill_chunk`. + # We use `tmp_end_idx` to store the end index of the kv cache to send. + self.tmp_end_idx: int = -1 + self.metadata_buffer_index: int = -1 + # Only meaningful in speculative reasoning. + self.accept_draft_tokens: float | None = None + + self.output_extra_info: dict[str, Any] = {} + + def set_tokenizer(self, tokenizer): + """Assign tokenizer and cache ids needed by check_finished().""" + self.tokenizer = tokenizer + if tokenizer is None: + self._eos_token_id_cached = None + self._additional_stop_token_ids_cached = None + return + eos_id = getattr(tokenizer, "eos_token_id", None) + self._eos_token_id_cached = int(eos_id) if eos_id is not None else None + extra = getattr(tokenizer, "additional_stop_token_ids", None) + self._additional_stop_token_ids_cached = ( + set(int(x) for x in extra) if extra else None + ) + + @property + def seqlen(self): + return len(self.origin_input_ids) + len(self.output_ids) + + def finished(self) -> bool: + # Whether request reached finished condition + return self.finished_reason is not None + + def init_incremental_detokenize(self): + first_iter = self.surr_offset is None or self.read_offset is None + + if first_iter: + self.read_offset = len(self.origin_input_ids_unpadded) + self.surr_offset = max( + self.read_offset - INIT_INCREMENTAL_DETOKENIZATION_OFFSET, 0 + ) + # self.surr_offset = self.read_offset + + all_ids = self.origin_input_ids_unpadded + self.output_ids + return all_ids[self.surr_offset :], self.read_offset - self.surr_offset + + def check_finished(self): + if self.finished(): + return + + if self.to_abort: + self.finished_reason = FINISH_ABORT( + message=self.to_abort_message, + ) + return + + if len(self.output_ids) >= self.sampling_params.max_new_tokens: + self.finished_reason = FINISH_LENGTH( + length=self.sampling_params.max_new_tokens + ) + return + + if self.grammar is not None: + if self.grammar.is_terminated(): + self.finished_reason = FINISH_MATCHED_TOKEN(matched=self.output_ids[-1]) + return + + last_token_id = self.output_ids[-1] + + if not self.sampling_params.ignore_eos: + matched_eos = False + + # Check stop token ids + if self.sampling_params.stop_token_ids: + matched_eos = last_token_id in self.sampling_params.stop_token_ids + if self.eos_token_ids: + matched_eos |= last_token_id in self.eos_token_ids + if self.tokenizer is not None and self._eos_token_id_cached is None: + self.set_tokenizer(self.tokenizer) + if self._eos_token_id_cached is not None: + matched_eos |= last_token_id == self._eos_token_id_cached + if self._additional_stop_token_ids_cached: + matched_eos |= last_token_id in self._additional_stop_token_ids_cached + if matched_eos: + self.finished_reason = FINISH_MATCHED_TOKEN(matched=last_token_id) + return + + # Check stop strings + if len(self.sampling_params.stop_strs) > 0: + tail_str = self.tokenizer.decode( + self.output_ids[-(self.sampling_params.stop_str_max_len + 1) :] + ) + + for stop_str in self.sampling_params.stop_strs: + if stop_str in tail_str or stop_str in self.decoded_text: + self.finished_reason = FINISH_MATCHED_STR(matched=stop_str) + return + + def __repr__(self): + return ( + f"Req(rid={self.rid}, " + f"input_ids={len(self.origin_input_ids)}, output_ids={len(self.output_ids)})" + ) diff --git a/python/tokenspeed/runtime/engine/request_handler.py b/python/tokenspeed/runtime/engine/request_handler.py new file mode 100644 index 0000000..ae2515c --- /dev/null +++ b/python/tokenspeed/runtime/engine/request_handler.py @@ -0,0 +1,500 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import logging +import os +from pathlib import Path +from typing import TYPE_CHECKING + +import torch +import zmq +from viztracer import VizTracer + +from tokenspeed.runtime.distributed.process_group_manager import ( + process_group_manager as pg_manager, +) +from tokenspeed.runtime.engine.generation_output_processor import RequestState +from tokenspeed.runtime.engine.io_struct import ( + AbortReq, + FlushCacheReqInput, + FlushCacheReqOutput, + GetInternalStateReq, + GetInternalStateReqOutput, + GetLoadReqInput, + GetLoadReqOutput, + IsSchedulerPausedReqInput, + IsSleepingReqInput, + PauseSchedulerReqInput, + ProfileReq, + ProfileReqOutput, + ProfileReqType, + ReleaseMemoryOccupationReqInput, + ResumeMemoryOccupationReqInput, + ResumeSchedulerReqInput, + SetInternalStateReq, + SetInternalStateReqOutput, + TokenizedGenerateReqInput, +) +from tokenspeed.runtime.engine.request_types import FINISH_ABORT +from tokenspeed.runtime.engine.scheduler_utils import make_spec +from tokenspeed.runtime.execution.forward_batch_info import ForwardMode +from tokenspeed.runtime.grammar.grammar_manager import GrammarManager +from tokenspeed.runtime.multimodal.shm_transport import sync_shm_features +from tokenspeed.runtime.pd.base.bootstrap import BootstrapInfo +from tokenspeed.runtime.utils import broadcast_pyobj +from tokenspeed.runtime.utils.dispatch import TypeBasedDispatcher +from tokenspeed.runtime.utils.env import envs +from tokenspeed.runtime.utils.hf_transformers_utils import get_tokenizer + +if TYPE_CHECKING: + from tokenspeed.runtime.utils.server_args import ServerArgs + +logger = logging.getLogger(__name__) + + +class RequestHandler: + """ + 1. Recv Reqs from ZMQ + 2. manage sessions + """ + + def __init__( + self, + server_args: ServerArgs, + hf_eos_token_id, + max_req_len: int, + vocab_size: int, + recv_func, + send_func, + get_load_fn=None, + architectures: list[str] | None = None, + pause_controller=None, + memory_controller=None, + ) -> None: + + self.forward_ct = 0 + self.server_args = server_args + # Owns pause/resume state; shared with the event loop. See pause.py. + self.pause_controller = pause_controller + # Owns release/resume_memory_occupation (data plane). See + # memory_occupation.py. Shares the pause controller's drain machinery. + self.memory_controller = memory_controller + + mapping = server_args.mapping + self.attn_tp_size = mapping.attn.tp_size + self.attn_tp_rank = mapping.attn.tp_rank + self.attn_tp_cpu_group = pg_manager.get_process_group( + "gloo", mapping.attn.tp_group + ) + self.attn_tp_src_rank = mapping.attn.tp_group[0] + + self.hf_eos_token_id = hf_eos_token_id + self.max_req_len = max_req_len + self.vocab_size = vocab_size + self.get_load_fn = get_load_fn + + self.tokenizer = get_tokenizer( + server_args.tokenizer, + tokenizer_mode=server_args.tokenizer_mode, + trust_remote_code=server_args.trust_remote_code, + revision=server_args.revision, + architectures=architectures, + ) + + self.recv_func = recv_func + self.send_func = send_func + + self.control_request_dispatcher = TypeBasedDispatcher( + [(ProfileReq, self.profile)] + ) + + self.grammar_manager = GrammarManager( + self.server_args, self.tokenizer, self.vocab_size + ) + + self.init_profiler() + + def recv_reqs(self) -> list: + """Receive results at attn_tp_rank = 0 and broadcast it to all other TP ranks.""" + if self.attn_tp_rank == 0: + recv_reqs = [] + + while True: + try: + recv_req = self.recv_func.recv_pyobj(zmq.NOBLOCK) + except zmq.ZMQError: + break + recv_reqs.append(recv_req) + else: + recv_reqs = None + + if self.attn_tp_size != 1: + recv_reqs = broadcast_pyobj( + recv_reqs, + self.attn_tp_rank, + self.attn_tp_cpu_group, + src=self.attn_tp_src_rank, + ) + + if recv_reqs: + sync_shm_features(recv_reqs, self.attn_tp_cpu_group, self.attn_tp_size) + + return recv_reqs + + def process_requests(self, recv_reqs: list): + """Dispatch control requests and return new generate request specs and states.""" + new_req_specs, req_states, bootstrap_infos, abort_rids = [], [], [], [] + for recv_req in recv_reqs: + if isinstance(recv_req, TokenizedGenerateReqInput): + req_spec, req_state, bootstrap_info = self.handle_generate_request( + recv_req + ) + + new_req_specs.append(req_spec) + req_states.append(req_state) + bootstrap_infos.append(bootstrap_info) + elif isinstance(recv_req, ProfileReq): + output = self.control_request_dispatcher(recv_req) + if output is not None: + self.send_func.send_pyobj(output) + elif isinstance(recv_req, AbortReq): + logger.debug("AbortReq for rid=%s", recv_req.rid) + abort_rids.append(recv_req.rid) + elif isinstance(recv_req, FlushCacheReqInput): + # Prefix cache is owned by the scheduler path; acknowledge the + # control request here so API callers still get a typed reply. + self.send_func.send_pyobj(FlushCacheReqOutput(success=True)) + elif isinstance(recv_req, PauseSchedulerReqInput): + # State change + reply (abort/wait replies are deferred by the + # controller until the event loop observes a drained scheduler). + self.pause_controller.handle_pause(recv_req) + elif isinstance(recv_req, ResumeSchedulerReqInput): + self.pause_controller.handle_resume(recv_req) + elif isinstance(recv_req, IsSchedulerPausedReqInput): + self.pause_controller.handle_is_paused(recv_req) + elif isinstance(recv_req, ReleaseMemoryOccupationReqInput): + # Deferred: pauses + drains, then frees GPU memory and replies. + self.memory_controller.handle_release(recv_req) + elif isinstance(recv_req, ResumeMemoryOccupationReqInput): + self.memory_controller.handle_resume(recv_req) + elif isinstance(recv_req, IsSleepingReqInput): + self.memory_controller.handle_is_sleeping(recv_req) + elif isinstance(recv_req, GetInternalStateReq): + self.send_func.send_pyobj(GetInternalStateReqOutput(internal_state={})) + elif isinstance(recv_req, SetInternalStateReq): + self.send_func.send_pyobj( + SetInternalStateReqOutput(updated=False, server_args={}) + ) + elif isinstance(recv_req, GetLoadReqInput): + if self.get_load_fn is not None: + self.send_func.send_pyobj(self.get_load_fn()) + else: + self.send_func.send_pyobj(GetLoadReqOutput()) + else: + raise NotImplementedError(f"Unsupported request type: {type(recv_req)}") + return new_req_specs, req_states, bootstrap_infos, abort_rids + + def handle_generate_request( + self, + recv_req: TokenizedGenerateReqInput, + ): + if recv_req.bootstrap_port is None: + recv_req.bootstrap_port = self.server_args.disaggregation_bootstrap_port + + req_spec = make_spec( + rid=recv_req.rid, + tokens=recv_req.input_ids, + ) + req_state = RequestState.from_recv_req( + recv_req, + tokenizer=self.tokenizer, + eos_token_ids=self.hf_eos_token_id, + ) + + if ( + recv_req.session_params is not None + and recv_req.session_params.id is not None + ): + req_state.finished_reason = FINISH_ABORT( + f"Invalid request: session id {recv_req.session_params.id} does not exist" + ) + return ( + req_spec, + req_state, + BootstrapInfo( + recv_req.bootstrap_host, + recv_req.bootstrap_port, + recv_req.bootstrap_room, + ), + ) + + req_state.sampling_params.max_new_tokens = min( + ( + req_state.sampling_params.max_new_tokens + if req_state.sampling_params.max_new_tokens is not None + else 1 << 30 + ), + self.max_req_len - len(req_state.prompt_input_ids) - 1, + ) + return ( + req_spec, + req_state, + BootstrapInfo( + recv_req.bootstrap_host, + recv_req.bootstrap_port, + recv_req.bootstrap_room, + ), + ) + + # ------------------------------------------------------------------ + # Profiling: torch / cuda / viztracer / mem-snapshot, driven by + # /start_profile and /stop_profile control requests. + # ------------------------------------------------------------------ + + def init_profiler(self): + self.torch_profiler = None + self.profiler_output_dir: str | None = None + self.profiler_activities: list[str] | None = None + self.profile_id: str | None = None + self.profiler_start_forward_ct: int | None = None + self.profiler_target_forward_ct: int | None = None + self.profiler_target_prefill_ct: int | None = None + self.profiler_target_decode_ct: int | None = None + self.profiler_prefill_ct: int | None = None + self.profiler_decode_ct: int | None = None + self.profile_by_stage: bool = False + self.profile_in_progress: bool = False + self.viztracer = None + + def init_profile( + self, + output_dir: str | None, + start_step: int | None, + num_steps: int | None, + activities: list[str] | None, + with_stack: bool | None, + record_shapes: bool | None, + profile_by_stage: bool, + profile_id: str, + ) -> ProfileReqOutput: + if self.profile_in_progress: + return ProfileReqOutput( + success=False, + message="Profiling is already in progress. Call /stop_profile first.", + ) + + self.profile_by_stage = profile_by_stage + + if output_dir is None: + output_dir = envs.TOKENSPEED_PROFILER_DIR.get() + if activities is None: + activities = ["CPU", "GPU"] + + self.profiler_output_dir = output_dir + self.torch_profiler_with_stack = with_stack + self.torch_profiler_record_shapes = record_shapes + self.profiler_activities = activities + self.profile_id = profile_id + + if start_step: + self.profiler_start_forward_ct = max(start_step, self.forward_ct + 1) + + if num_steps: + if self.profile_by_stage: + self.profiler_target_prefill_ct = num_steps + self.profiler_target_decode_ct = num_steps + self.profiler_prefill_ct = 0 + self.profiler_decode_ct = 0 + elif start_step: + self.profiler_target_forward_ct = ( + self.profiler_start_forward_ct + num_steps + ) + else: + self.profiler_target_forward_ct = self.forward_ct + num_steps + # The caller will be notified when reaching profiler_target_forward_ct + else: + self.profiler_target_forward_ct = None + + return ProfileReqOutput(success=True, message="Succeeded") + + def start_profile( + self, stage: ForwardMode | None = None + ) -> ProfileReqOutput | None: + stage_str = f" for {stage.name}" if stage else "" + stage_suffix = f"-{stage.name}" if stage else "" + + activities = self.profiler_activities + with_stack = self.torch_profiler_with_stack + record_shapes = self.torch_profiler_record_shapes + + activity_map = { + "CPU": torch.profiler.ProfilerActivity.CPU, + "GPU": torch.profiler.ProfilerActivity.CUDA, + } + torchprof_activities = [ + activity_map[a] for a in activities if a in activity_map + ] + + if torchprof_activities: + self.torch_profiler = torch.profiler.profile( + activities=torchprof_activities, + with_stack=with_stack if with_stack is not None else True, + record_shapes=record_shapes if record_shapes is not None else False, + ) + self.torch_profiler.start() + + if "MEM" in activities: + torch.cuda.memory._record_memory_history(max_entries=100000) + + if "CUDA_PROFILER" in activities: + torch.cuda.cudart().cudaProfilerStart() + + if "VIZTRACER" in activities: + Path(self.profiler_output_dir).mkdir(parents=True, exist_ok=True) + self.viztracer = VizTracer( + output_file=os.path.join( + self.profiler_output_dir, + f"{self.profile_id}-TP-{self.attn_tp_rank}{stage_suffix}.viztracer.json", + ), + min_duration=int( + os.environ.get("TOKENSPEED_VIZTRACER_MIN_DURATION_US", "100") + ), + log_async=True, + ) + self.viztracer.start() + + if activities: + if activities != ["CUDA_PROFILER"]: + logger.info( + "Profiling starts%s. Traces will be saved to: %s (with profile id: %s)", + stage_str, + self.profiler_output_dir, + self.profile_id, + ) + self.profile_in_progress = True + + return ProfileReqOutput(success=True, message="Succeeded") + + def stop_profile(self, stage: ForwardMode | None = None) -> ProfileReqOutput | None: + if not self.profile_in_progress: + return ProfileReqOutput( + success=False, + message="Profiling is not in progress. Call /start_profile first.", + ) + + Path(self.profiler_output_dir).mkdir(parents=True, exist_ok=True) + + stage_suffix = f"-{stage.name}" if stage else "" + logger.info("Stop profiling%s...", stage_suffix) + + if self.torch_profiler is not None: + self.torch_profiler.stop() + self.torch_profiler.export_chrome_trace( + os.path.join( + self.profiler_output_dir, + f"{self.profile_id}-TP-{self.attn_tp_rank}{stage_suffix}.trace.json.gz", + ) + ) + torch.distributed.barrier(self.attn_tp_cpu_group) + + if self.profiler_activities is not None and "MEM" in self.profiler_activities: + memory_profile_path = os.path.join( + self.profiler_output_dir, + f"{self.profile_id}-TP-{self.attn_tp_rank}-memory{stage_suffix}.pickle", + ) + torch.cuda.memory._dump_snapshot(memory_profile_path) + torch.cuda.memory._record_memory_history(enabled=None) + + if "CUDA_PROFILER" in self.profiler_activities: + torch.cuda.cudart().cudaProfilerStop() + + if "VIZTRACER" in self.profiler_activities and self.viztracer is not None: + self.viztracer.stop() + self.viztracer.save() + self.viztracer = None + + if self.profiler_activities and self.profiler_activities != ["CUDA_PROFILER"]: + logger.info( + "Profiling done. Traces are saved to: %s", self.profiler_output_dir + ) + + self.torch_profiler = None + self.profile_in_progress = False + self.profiler_start_forward_ct = None + + return ProfileReqOutput(success=True, message="Succeeded.") + + def _profile_batch_predicate(self, forward_mode=None): + """Check and toggle profiling based on forward step count. + + Args: + forward_mode: Optional ForwardMode for stage-based profiling. + Not needed for step-count-based profiling. + """ + if self.profile_by_stage and forward_mode is not None: + if forward_mode.is_extend_or_mixed(): + if self.profiler_prefill_ct == 0: + self.start_profile(forward_mode) + self.profiler_prefill_ct += 1 + if self.profiler_prefill_ct > self.profiler_target_prefill_ct: + if self.profile_in_progress: + self.stop_profile(stage=ForwardMode.EXTEND) + elif forward_mode.is_decode(): + if self.profiler_decode_ct == 0: + if self.profile_in_progress: + self.stop_profile(ForwardMode.EXTEND) + self.start_profile(forward_mode) + self.profiler_decode_ct += 1 + if self.profiler_decode_ct > self.profiler_target_decode_ct: + if self.profile_in_progress: + self.stop_profile(stage=ForwardMode.DECODE) + elif forward_mode.is_idle(): + pass + else: + if ( + self.profiler_target_forward_ct + and self.profiler_target_forward_ct <= self.forward_ct + ): + self.stop_profile() + if ( + self.profiler_start_forward_ct + and self.profiler_start_forward_ct == self.forward_ct + ): + self.start_profile() + + def profile(self, recv_req: ProfileReq): + if recv_req.type == ProfileReqType.START_PROFILE: + res = self.init_profile( + recv_req.output_dir, + recv_req.start_step, + recv_req.num_steps, + recv_req.activities, + recv_req.with_stack, + recv_req.record_shapes, + recv_req.profile_by_stage, + recv_req.profile_id, + ) + if not res.success or recv_req.profile_by_stage or recv_req.start_step: + return res + return self.start_profile() + else: + return self.stop_profile() diff --git a/python/tokenspeed/runtime/engine/request_stats.py b/python/tokenspeed/runtime/engine/request_stats.py new file mode 100644 index 0000000..168b879 --- /dev/null +++ b/python/tokenspeed/runtime/engine/request_stats.py @@ -0,0 +1,218 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Per-request performance stats for --enable-log-request-stats. + +Two cohesive pieces: +- ``RequestStatsTracker``: a mutable, host-side accumulator that the engine fires + lifecycle events into (scheduled/prefill-done/first-token/finish, preemption). + ``NOOP_STATS`` is a shared null-object instance carried by requests until + ``register()`` attaches a real tracker, so call sites need no per-request guard. +- ``RequestStats``: the immutable, rounded summary derived from a tracker + + RequestState via ``from_state``, logged as a Python-object repr. + +Everything here is host-side; building these introduces no GPU sync. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from tokenspeed.runtime.engine.request_types import FINISH_ABORT + +if TYPE_CHECKING: + from tokenspeed.runtime.engine.generation_output_processor import RequestState + + +def _ms(end: float, start: float) -> float: + # 0 for unset (0.0) timestamps, not a garbage epoch-sized number + return round((end - start) * 1e3, 2) if end > 0.0 and start > 0.0 else 0.0 + + +class RequestStatsTracker: + """Host-side per-request timing/preemption accumulator (--enable-log-request-stats). + + Attached to a RequestState only when logging is on; the engine fires + lifecycle events into it. No GPU sync. + """ + + def __init__(self) -> None: + self.scheduled_time = 0.0 + self.prefill_done_time = 0.0 + self.first_token_time = 0.0 + self.finish_time = 0.0 + self.preempt_count = 0 + self.preempt_time = 0.0 + self._preempted_last_step = False + + def mark_scheduled(self, now: float) -> None: + if self.scheduled_time == 0.0: + self.scheduled_time = now + + def mark_prefill_done(self, now: float) -> None: + if self.prefill_done_time == 0.0: + self.prefill_done_time = now + + def mark_first_token(self, now: float) -> None: + if self.first_token_time == 0.0: + self.first_token_time = now + + def mark_finish(self, now: float) -> None: + self.finish_time = now + + def record_decode_step(self, step_dt: float, prefilling_others: bool) -> None: + # count each prefill-of-others interruption (rising edge) and its time + if prefilling_others: + if not self._preempted_last_step: + self.preempt_count += 1 + self.preempt_time += step_dt + self._preempted_last_step = True + else: + self._preempted_last_step = False + + +class _NoOpStatsTracker(RequestStatsTracker): + """Null-object tracker used when --enable-log-request-stats is off, so the engine + can fire lifecycle events unconditionally without a per-request guard. The + only stats check lives in _log_request_stats (which skips this singleton). + + Frozen on purpose: it is a SHARED singleton, so any attribute write raises. + If a future tracker mutator is added without a no-op override here, the call + falls through to the base method and fails loudly instead of silently + corrupting shared state across all flag-off requests. + """ + + def __init__(self) -> None: + pass # carries no per-instance state; every method below is a no-op + + def __setattr__(self, name: str, value: object) -> None: + raise AttributeError("NOOP_STATS is a read-only no-op singleton") + + def mark_scheduled(self, now: float) -> None: + pass + + def mark_prefill_done(self, now: float) -> None: + pass + + def mark_first_token(self, now: float) -> None: + pass + + def mark_finish(self, now: float) -> None: + pass + + def record_decode_step(self, step_dt: float, prefilling_others: bool) -> None: + pass + + +# Shared singleton: requests carry this until register() attaches a real tracker. +NOOP_STATS = _NoOpStatsTracker() + + +@dataclass +class RequestStats: + """Host-side per-request perf summary (--enable-log-request-stats), logged as repr. + + Durations are ms, rates are 0-1, ``*_ts`` are epoch seconds; ``acc_*`` are + None when spec decode is off. No GPU sync. + """ + + status: str + reason: str + + prompt_tokens: int + cache_tokens: int + output_tokens: int + cache_hit_rate: float + + queue_ms: float + prefill_ms: float + ttft_ms: float + total_ms: float + preempt_ms: float + preempt_count: int + decode_tps: float + + acc_len: float | None + acc_rate: float | None + + recv_ts: float + commit_ts: float + finish_ts: float + + @classmethod + def from_state( + cls, rs: RequestState, spec_algorithm, spec_num_tokens + ) -> RequestStats: + t = rs.stats + prompt = rs.input_length + output = rs.output_length + + decode_window = ( + t.finish_time - t.first_token_time + if t.finish_time > 0.0 and t.first_token_time > 0.0 + else 0.0 + ) + decode_tps = ( + round((output - 1) / decode_window, 1) + if decode_window > 0.0 and output > 1 + else 0.0 + ) + + # spec acceptance; None when spec decode is off + if spec_algorithm is not None and rs.spec_verify_ct > 0: + acc_len = rs.accept_draft_tokens or 0.0 + acc_rate = ( + round(max(0.0, acc_len - 1.0) / spec_num_tokens, 4) + if spec_num_tokens + else 0.0 + ) + acc_len = round(acc_len, 2) + else: + acc_len = acc_rate = None + + return cls( + status=( + "aborted" + if isinstance(rs.finished_reason, FINISH_ABORT) + else "finished" + ), + reason=( + rs.finished_reason.to_json().get("type", "unknown") + if rs.finished_reason is not None + else "unknown" + ), + prompt_tokens=prompt, + cache_tokens=rs.cached_tokens, + output_tokens=output, + cache_hit_rate=round(rs.cached_tokens / prompt, 4) if prompt > 0 else 0.0, + queue_ms=_ms(t.scheduled_time, rs.created_time), + prefill_ms=_ms(t.prefill_done_time, t.scheduled_time), + ttft_ms=_ms(t.first_token_time, rs.created_time), + total_ms=_ms(t.finish_time, rs.created_time), + preempt_ms=round(t.preempt_time * 1e3, 2), + preempt_count=t.preempt_count, + decode_tps=decode_tps, + acc_len=acc_len, + acc_rate=acc_rate, + recv_ts=round(rs.created_time, 3), + commit_ts=round(t.scheduled_time, 3), + finish_ts=round(t.finish_time, 3), + ) diff --git a/python/tokenspeed/runtime/engine/request_types.py b/python/tokenspeed/runtime/engine/request_types.py new file mode 100644 index 0000000..eca2d3a --- /dev/null +++ b/python/tokenspeed/runtime/engine/request_types.py @@ -0,0 +1,95 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Shared request types used by both scheduler and output processing.""" + +from enum import Enum + +INIT_INCREMENTAL_DETOKENIZATION_OFFSET = 5 + + +class BaseFinishReason: + def __init__(self, is_error: bool = False): + self.is_error = is_error + + def to_json(self): + raise NotImplementedError() + + +class FINISH_MATCHED_TOKEN(BaseFinishReason): + def __init__(self, matched: int | list[int]): + super().__init__() + self.matched = matched + + def to_json(self): + return { + "type": "stop", + "matched": self.matched, + } + + +class FINISH_MATCHED_STR(BaseFinishReason): + def __init__(self, matched: str): + super().__init__() + self.matched = matched + + def to_json(self): + return { + "type": "stop", + "matched": self.matched, + } + + +class FINISH_LENGTH(BaseFinishReason): + def __init__(self, length: int): + super().__init__() + self.length = length + + def to_json(self): + return { + "type": "length", + "length": self.length, + } + + +class ABORT_CODE(Enum): + TransferFailed = 521 + UnknownError = 522 + NumericalError = 523 + + +class FINISH_ABORT(BaseFinishReason): + def __init__( + self, + message: str = "Unknown error", + status_code=None, + err_type=ABORT_CODE.UnknownError, + ): + super().__init__(is_error=True) + self.message = message + self.status_code = status_code + self.err_type = err_type + + def to_json(self): + return { + "type": "abort", + "message": self.message, + "err_type": self.err_type.value, + } diff --git a/python/tokenspeed/runtime/engine/schedule_batch.py b/python/tokenspeed/runtime/engine/schedule_batch.py new file mode 100755 index 0000000..131b30b --- /dev/null +++ b/python/tokenspeed/runtime/engine/schedule_batch.py @@ -0,0 +1,332 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +""" +Store information about requests and batches. + +The following is the flow of data structures for a batch: + +ScheduleBatch -> executor inputs + +- ScheduleBatch is managed by the runtime event loop and model executor. + It contains high-level scheduling data. Most of the data is on the CPU. +- Executor inputs contain low-level tensor data. Most of the data consists of + GPU tensors. +""" + +from __future__ import annotations + +import dataclasses +import threading +from collections.abc import Callable +from typing import TYPE_CHECKING + +import torch +import triton +import triton.language as tl + +from tokenspeed.runtime.cache.allocator import KVAllocator +from tokenspeed.runtime.cache.base_prefix_cache import BasePrefixCache +from tokenspeed.runtime.cache.req_to_token_pool import ReqToTokenPool +from tokenspeed.runtime.configs.model_config import ModelConfig +from tokenspeed.runtime.engine.request import Req +from tokenspeed.runtime.execution.forward_batch_info import ( + ForwardMode, +) +from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool +from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput +from tokenspeed.runtime.pd.disaggregation_decode_scheduler import ( + DisaggDecodeScheduler, +) +from tokenspeed.runtime.sampling.sampling_batch_info import SamplingBatchInfo +from tokenspeed.runtime.utils import get_colorful_logger + +if TYPE_CHECKING: + from tokenspeed.runtime.spec_decode.algorithm import SpeculativeAlgorithm + from tokenspeed.runtime.spec_decode.eagle import EagleDraftInput + + +logger = get_colorful_logger(__name__) + +bid = 0 + + +@dataclasses.dataclass +class ScheduleBatch(DisaggDecodeScheduler): + """Store all information of a batch on the scheduler.""" + + # Request, memory pool, and cache + reqs: list[Req] + req_to_token_pool: ReqToTokenPool = None + kv_allocator: KVAllocator = None + token_to_kv_pool: BaseTokenToKVPool = None + tree_cache: BasePrefixCache = None + + # Batch configs + model_config: ModelConfig = None + forward_mode: ForwardMode = None + enable_overlap: bool = False + + # Events + launch_done: threading.Event | None = None + + # Sampling info + sampling_info: SamplingBatchInfo = None + next_batch_sampling_info: SamplingBatchInfo = None + + # Batched arguments to model runner + input_ids: torch.Tensor = None # shape: [b], int32 + input_multi_ids: torch.Tensor | None = None # shape: [b, mm_heads], int32 + draft_input_ids: torch.Tensor = None # shape: [b], int32 + input_embeds: torch.Tensor = None # shape: [b, hidden_size], float32 + input_extra_infos: list[dict] | None = None + req_pool_indices: torch.Tensor = None # shape: [b], int32 + + seq_lens: torch.Tensor = None # shape: [b], int64 + output_ids: torch.Tensor = None # shape: [b], int32 + output_multi_ids: torch.Tensor = None # shape: [b], int32 + + # The sum of all sequence lengths + seq_lens_sum: int = None + + # For DP attention + global_num_tokens: list[int] | None = ( + None # e.g. dp = 4, attn-tp = 2, [A, A, B, B, C, C, D, D] + ) + global_num_tokens_for_logprob: list[int] | None = None + all_decode_or_idle: bool = False + # For processing logprobs + return_logprob: bool = False + top_logprobs_nums: list[int] | None = None + token_ids_logprobs: list[list[int]] | None = None + + # For extend and mixed chunekd prefill + prefix_lens: list[int] = None + extend_lens: list[int] = None + extend_num_tokens: int = None + decoding_reqs: list[Req] = None + extend_logprob_start_lens: list[int] = None + # It comes empty list if logprob is not required. + extend_input_logprob_token_ids: torch.Tensor | None = None + + # Stream + has_stream: bool = False + + # Has grammar + has_grammar: bool = False + + # Device + device: str = "cuda" + + # Speculative decoding + spec_algorithm: SpeculativeAlgorithm = None + spec_info: EagleDraftInput | None = None + draft_token_num: int | None = 0 + spec_num_steps: int | None = 0 + # Reserve multiple positions for speculative decoding + reserve_num_tokens_init: int = None + + # Enable custom logit processor + enable_custom_logit_processor: bool = False + + # Whether to return hidden states + return_hidden_states: bool = False + + # set aux data for Disaggregation + disagg_set_aux_fn: Callable[[torch.Tensor, LogitsProcessorOutput], None] | None = ( + None + ) + # kvstore pointer for synchronizing data loading from CPU to GPU + kvstore_consumer_index: int = -1 + + @classmethod + def init_new( + cls, + reqs: list[Req], + req_to_token_pool: ReqToTokenPool, + kv_allocator: KVAllocator, + token_to_kv_pool: BaseTokenToKVPool, + tree_cache: BasePrefixCache, + model_config: ModelConfig, + enable_overlap: bool, + spec_algorithm: SpeculativeAlgorithm, + enable_custom_logit_processor: bool, + reserve_num_tokens_init: int = 0, + draft_token_num: int = 0, + spec_num_steps: int = 0, + ): + return cls( + reqs=reqs, + req_to_token_pool=req_to_token_pool, + kv_allocator=kv_allocator, + token_to_kv_pool=token_to_kv_pool, + tree_cache=tree_cache, + model_config=model_config, + enable_overlap=enable_overlap, + return_logprob=any(req.return_logprob for req in reqs), + has_stream=any(req.stream for req in reqs), + has_grammar=any(req.grammar for req in reqs), + device=req_to_token_pool.device, + spec_algorithm=spec_algorithm, + enable_custom_logit_processor=enable_custom_logit_processor, + return_hidden_states=any(req.return_hidden_states for req in reqs), + reserve_num_tokens_init=reserve_num_tokens_init, + draft_token_num=draft_token_num, + spec_num_steps=spec_num_steps, + ) + + def batch_size(self): + return len(self.reqs) + + def alloc_token_slots(self, req_pool_index: int, num_tokens: int): + out_cache_loc = self.kv_allocator.alloc( + req_pool_index, + num_tokens, + self.req_to_token_pool.alloced_lens[req_pool_index].item(), + ) + + if out_cache_loc is None: + if self.tree_cache is not None: + logger.debug( + "[evict] before evict evict_tokens=%s evictable_size=%s", + num_tokens, + self.tree_cache.evictable_size(), + ) + need_page_num = ( + num_tokens + self.kv_allocator.page_size - 1 + ) // self.kv_allocator.page_size + self.tree_cache.evict(need_page_num, self.kv_allocator.free) + logger.debug( + "[evict] after evict evictable_size=%s", + self.tree_cache.evictable_size(), + ) + out_cache_loc = self.kv_allocator.alloc( + req_pool_index, + num_tokens, + self.req_to_token_pool.alloced_lens[req_pool_index].item(), + ) + logger.debug("[evict] out_cache_loc=%r after evict", out_cache_loc) + + if out_cache_loc is None: + phase_str = ( + "Prefill" if self.forward_mode.is_extend_or_mixed() else "Decode" + ) + logger.error( + "%s out of memory. Try to lower your batch size.\nTry to allocate %s tokens.\nAvailable tokens: %s\n", + phase_str, + num_tokens, + self.kv_allocator.available_size() + + self.tree_cache.evictable_size(), + ) + if self.tree_cache is not None: + self.tree_cache.pretty_print() + exit(1) + + return out_cache_loc + + def prealloc_for_draft_decode(self, is_disaggregation_decode: bool = False): + """Pre-allocate a segment of slots for draft decode""" + if self.enable_overlap: + # Conceptually, each allocation during speculation + overlap is preparing for the next batch's launch. + # Therefore, at the beginning, reserve enough space at the end of prefill for the next round's verify and draft decode. + # Then, each time adjust the reserved space based on acceptance length to prevent allocation divergence causing insufficient space. + # The reserved space for draft decode will always be overwritten by valid tokens in the next verify. + # Initially allocate spec_num_steps, subsequent allocations are not needed. + num_tokens_pre_alloc = self.draft_token_num + (self.spec_num_steps - 1) + else: + # Synchronously, each allocation is for the current batch's launch. Here we allocate spec_num_steps + # extra slots to reserve enough space for draft decode. + if self.spec_num_steps > 1: + num_tokens_pre_alloc = self.spec_num_steps - 1 + else: + return + out_cache_loc_list = [] + req_indices = [] + for i, req in enumerate(self.reqs): + # End of prefill or PD disaggregation mocked prefill + if req.draft_fill_ids[-1] == -1 or is_disaggregation_decode: + out_cache_loc_list.append( + self.alloc_token_slots(req.req_pool_idx, num_tokens_pre_alloc) + ) + req_indices.append(req.req_pool_idx) + bs = len(req_indices) + if len(out_cache_loc_list) == 0: + return + out_cache_loc = torch.concat(out_cache_loc_list) + out_cache_loc = out_cache_loc.to(self.device, non_blocking=True) + req_indices = torch.tensor(req_indices, dtype=torch.int64).to( + self.device, non_blocking=True + ) + start_offsets = torch.index_select( + self.req_to_token_pool.alloced_lens, 0, req_indices + ) + end_offsets = start_offsets + num_tokens_pre_alloc + assign_req_to_token_pool[(bs,)]( + req_indices, + self.req_to_token_pool.req_to_token, + start_offsets, + end_offsets, + out_cache_loc, + self.req_to_token_pool.req_to_token.shape[1], + triton.next_power_of_2(bs), + ) + self.req_to_token_pool.alloced_lens[req_indices] += num_tokens_pre_alloc + + def __str__(self): + return ( + f"ScheduleBatch(forward_mode={self.forward_mode.name}, " + f"#req={(len(self.reqs))})" + ) + + +@triton.jit +def assign_req_to_token_pool( + req_pool_indices, + req_to_token, + start_offset, + end_offset, + out_cache_loc, + pool_len: tl.constexpr, + bs_upper: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 32 + pid = tl.program_id(axis=0) + kv_start = tl.load(start_offset + pid) + kv_end = tl.load(end_offset + pid) + token_pool = req_to_token + tl.load(req_pool_indices + pid) * pool_len + + # Get the offset for reading out_cache + length_offset = tl.arange(0, bs_upper) + start = tl.load(start_offset + length_offset, mask=length_offset < pid) + end = tl.load(end_offset + length_offset, mask=length_offset < pid) + out_offset = tl.sum(end - start, axis=0) + + out_cache_ptr = out_cache_loc + out_offset + + save_offset = tl.arange(0, BLOCK_SIZE) + kv_start + load_offset = tl.arange(0, BLOCK_SIZE) + + num_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE) + for _ in range(num_loop): + mask = save_offset < kv_end + data = tl.load(out_cache_ptr + load_offset, mask=mask) + tl.store(token_pool + save_offset, data, mask=mask) + save_offset += BLOCK_SIZE + load_offset += BLOCK_SIZE diff --git a/python/tokenspeed/runtime/engine/scheduler_control_client.py b/python/tokenspeed/runtime/engine/scheduler_control_client.py new file mode 100755 index 0000000..ee76162 --- /dev/null +++ b/python/tokenspeed/runtime/engine/scheduler_control_client.py @@ -0,0 +1,452 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import asyncio +import copy +import logging +import time +from collections import deque +from typing import ( + TYPE_CHECKING, + Any, + Generic, + TypeVar, +) + +import zmq + +from tokenspeed.runtime.engine.io_struct import ( + ExpertDistributionReq, + ExpertDistributionReqOutput, + FlushCacheReqInput, + FlushCacheReqOutput, + GetInternalStateReq, + GetInternalStateReqOutput, + GetLoadReqInput, + GetLoadReqOutput, + GetWeightsByNameReqInput, + GetWeightsByNameReqOutput, + InitWeightsUpdateGroupReqInput, + InitWeightsUpdateGroupReqOutput, + IsSchedulerPausedReqInput, + IsSchedulerPausedReqOutput, + IsSleepingReqInput, + IsSleepingReqOutput, + PauseMode, + PauseSchedulerReqInput, + PauseSchedulerReqOutput, + ProfileReq, + ProfileReqOutput, + ProfileReqType, + ReleaseMemoryOccupationReqInput, + ReleaseMemoryOccupationReqOutput, + ResumeMemoryOccupationReqInput, + ResumeMemoryOccupationReqOutput, + ResumeSchedulerReqInput, + ResumeSchedulerReqOutput, + SetInternalStateReq, + SetInternalStateReqOutput, + UpdateWeightsFromDistributedReqInput, + UpdateWeightsFromDistributedReqOutput, + UpdateWeightsFromTensorReqInput, + UpdateWeightsFromTensorReqOutput, +) +from tokenspeed.runtime.utils.dispatch import TypeBasedDispatcher +from tokenspeed.runtime.utils.env import envs +from tokenspeed.runtime.utils.server_args import ServerArgs + +if TYPE_CHECKING: + from tokenspeed.runtime.engine.async_llm import AsyncLLM + +T = TypeVar("T") + +logger = logging.getLogger(__name__) + + +class _Communicator(Generic[T]): + """Note: The communicator now only run up to 1 in-flight request at any time.""" + + def __init__(self, sender: zmq.Socket, fan_out: int, mode="queueing"): + self._sender = sender + self._fan_out = fan_out + self._mode = mode + self._result_event: asyncio.Event | None = None + self._result_values: list[T] | None = None + self._ready_queue: deque[asyncio.Future] = deque() + + if mode not in ("queueing", "watching"): + raise ValueError(f"Invalid communicator mode: {mode}") + + async def queueing_call(self, obj: T): + ready_event = asyncio.Event() + if self._result_event is not None or len(self._ready_queue) > 0: + self._ready_queue.append(ready_event) + await ready_event.wait() + if self._result_event is not None or self._result_values is not None: + raise RuntimeError("Communicator result state was not reset.") + + if obj: + self._sender.send_pyobj(obj) + + self._result_event = asyncio.Event() + self._result_values = [] + await self._result_event.wait() + result_values = self._result_values + self._result_event = self._result_values = None + + if len(self._ready_queue) > 0: + self._ready_queue.popleft().set() + + return result_values + + async def watching_call(self, obj): + if self._result_event is None: + if self._result_values is not None: + raise RuntimeError("Communicator result values were not reset.") + self._result_values = [] + self._result_event = asyncio.Event() + + if obj: + self._sender.send_pyobj(obj) + + await self._result_event.wait() + result_values = copy.deepcopy(self._result_values) + self._result_event = self._result_values = None + return result_values + + async def __call__(self, obj): + if self._mode == "queueing": + return await self.queueing_call(obj) + else: + return await self.watching_call(obj) + + def handle_recv(self, recv_obj: T): + self._result_values.append(recv_obj) + if len(self._result_values) == self._fan_out: + self._result_event.set() + + +class SchedulerControlClient: + """Scheduler control-plane client methods for AsyncLLM.""" + + def init_communicators(self: AsyncLLM, server_args: ServerArgs): + # Communicators + self.init_weights_update_group_communicator = _Communicator( + self.engine_core_client.send_to_scheduler, server_args.mapping.attn.dp_size + ) + self.update_weights_from_distributed_communicator = _Communicator( + self.engine_core_client.send_to_scheduler, server_args.mapping.attn.dp_size + ) + self.update_weights_from_tensor_communicator = _Communicator( + self.engine_core_client.send_to_scheduler, server_args.mapping.attn.dp_size + ) + self.get_weights_by_name_communicator = _Communicator( + self.engine_core_client.send_to_scheduler, server_args.mapping.attn.dp_size + ) + self.release_memory_occupation_communicator = _Communicator( + self.engine_core_client.send_to_scheduler, server_args.mapping.attn.dp_size + ) + self.resume_memory_occupation_communicator = _Communicator( + self.engine_core_client.send_to_scheduler, server_args.mapping.attn.dp_size + ) + self.flush_cache_communicator = _Communicator( + self.engine_core_client.send_to_scheduler, server_args.mapping.attn.dp_size + ) + self.pause_scheduler_communicator = _Communicator( + self.engine_core_client.send_to_scheduler, server_args.mapping.attn.dp_size + ) + self.resume_scheduler_communicator = _Communicator( + self.engine_core_client.send_to_scheduler, server_args.mapping.attn.dp_size + ) + self.is_scheduler_paused_communicator = _Communicator( + self.engine_core_client.send_to_scheduler, server_args.mapping.attn.dp_size + ) + self.is_sleeping_communicator = _Communicator( + self.engine_core_client.send_to_scheduler, server_args.mapping.attn.dp_size + ) + self.profile_communicator = _Communicator( + self.engine_core_client.send_to_scheduler, server_args.mapping.attn.dp_size + ) + self.get_internal_state_communicator = _Communicator( + self.engine_core_client.send_to_scheduler, server_args.mapping.attn.dp_size + ) + self.set_internal_state_communicator = _Communicator( + self.engine_core_client.send_to_scheduler, server_args.mapping.attn.dp_size + ) + + self.expert_distribution_communicator = _Communicator( + self.engine_core_client.send_to_scheduler, server_args.mapping.attn.dp_size + ) + self.get_load_communicator = _Communicator( + self.engine_core_client.send_to_scheduler, + server_args.mapping.attn.dp_size, + mode="watching", + ) + + self._result_dispatcher += self._get_communicator_dispatcher() + + def _get_communicator_dispatcher(self: AsyncLLM): + return TypeBasedDispatcher( + [ + ( + InitWeightsUpdateGroupReqOutput, + self.init_weights_update_group_communicator.handle_recv, + ), + ( + UpdateWeightsFromDistributedReqOutput, + self.update_weights_from_distributed_communicator.handle_recv, + ), + ( + UpdateWeightsFromTensorReqOutput, + self.update_weights_from_tensor_communicator.handle_recv, + ), + ( + GetWeightsByNameReqOutput, + self.get_weights_by_name_communicator.handle_recv, + ), + ( + ReleaseMemoryOccupationReqOutput, + self.release_memory_occupation_communicator.handle_recv, + ), + ( + ResumeMemoryOccupationReqOutput, + self.resume_memory_occupation_communicator.handle_recv, + ), + ( + FlushCacheReqOutput, + self.flush_cache_communicator.handle_recv, + ), + ( + PauseSchedulerReqOutput, + self.pause_scheduler_communicator.handle_recv, + ), + ( + ResumeSchedulerReqOutput, + self.resume_scheduler_communicator.handle_recv, + ), + ( + IsSchedulerPausedReqOutput, + self.is_scheduler_paused_communicator.handle_recv, + ), + ( + IsSleepingReqOutput, + self.is_sleeping_communicator.handle_recv, + ), + ( + ProfileReqOutput, + self.profile_communicator.handle_recv, + ), + ( + GetInternalStateReqOutput, + self.get_internal_state_communicator.handle_recv, + ), + ( + SetInternalStateReqOutput, + self.set_internal_state_communicator.handle_recv, + ), + ( + ExpertDistributionReqOutput, + self.expert_distribution_communicator.handle_recv, + ), + ( + GetLoadReqOutput, + self.get_load_communicator.handle_recv, + ), + ] + ) + + async def flush_cache(self: AsyncLLM) -> FlushCacheReqOutput: + return (await self.flush_cache_communicator(FlushCacheReqInput()))[0] + + async def pause_scheduler(self: AsyncLLM, *, mode: PauseMode = "abort") -> bool: + """Pause generation to allow model weight updates. + + ``mode`` controls in-flight requests: ``"abort"`` cancels them, + ``"wait"`` lets them finish, ``"keep"`` freezes them for ``/resume``. + For ``abort``/``wait`` the reply only returns once the scheduler has + drained, so on return no forward work is in flight. + + Cache invalidation after a weight swap is the weight-update op's + responsibility (``update_weights_*(flush_cache=...)``), not pause's. + """ + # Pause may be the very first call (e.g. weight swap before serving), + # so ensure the output-dispatch loop is running to receive the reply. + self.auto_create_handle_loop() + result = ( + await self.pause_scheduler_communicator(PauseSchedulerReqInput(mode=mode)) + )[0] + return result.success + + async def resume_scheduler(self: AsyncLLM) -> bool: + """Resume generation after :meth:`pause_scheduler`.""" + self.auto_create_handle_loop() + result = (await self.resume_scheduler_communicator(ResumeSchedulerReqInput()))[ + 0 + ] + return result.success + + async def is_scheduler_paused(self: AsyncLLM) -> bool: + """Return whether the scheduler is currently paused.""" + self.auto_create_handle_loop() + result = ( + await self.is_scheduler_paused_communicator(IsSchedulerPausedReqInput()) + )[0] + return result.is_paused + + async def start_profile( + self: AsyncLLM, + output_dir: str | None = None, + start_step: int | None = None, + num_steps: int | None = None, + activities: list[str] | None = None, + with_stack: bool | None = None, + record_shapes: bool | None = None, + profile_by_stage: bool = False, + profile_id: str | None = None, + ): + self.auto_create_handle_loop() + env_with_stack = envs.TOKENSPEED_PROFILE_WITH_STACK.get() + with_stack = False if with_stack is False or env_with_stack is False else True + req = ProfileReq( + type=ProfileReqType.START_PROFILE, + output_dir=output_dir, + start_step=start_step, + num_steps=num_steps, + activities=activities, + with_stack=with_stack, + record_shapes=record_shapes, + profile_by_stage=profile_by_stage, + profile_id=profile_id or time.strftime("%Y%m%d-%H%M%S"), + ) + return await self._execute_profile(req) + + async def stop_profile(self: AsyncLLM): + self.auto_create_handle_loop() + req = ProfileReq(type=ProfileReqType.STOP_PROFILE) + return await self._execute_profile(req) + + async def _execute_profile(self: AsyncLLM, req: ProfileReq): + result = (await self.profile_communicator(req))[0] + if not result.success: + raise RuntimeError(result.message) + return result + + async def start_expert_distribution_record(self: AsyncLLM): + self.auto_create_handle_loop() + await self.expert_distribution_communicator(ExpertDistributionReq.START_RECORD) + + async def stop_expert_distribution_record(self: AsyncLLM): + self.auto_create_handle_loop() + await self.expert_distribution_communicator(ExpertDistributionReq.STOP_RECORD) + + async def dump_expert_distribution_record(self: AsyncLLM): + self.auto_create_handle_loop() + await self.expert_distribution_communicator(ExpertDistributionReq.DUMP_RECORD) + + async def init_weights_update_group( + self: AsyncLLM, + obj: InitWeightsUpdateGroupReqInput, + ) -> tuple[bool, str]: + self.auto_create_handle_loop() + if self.server_args.mapping.attn.has_dp: + raise RuntimeError("dp_size must be 1 for init parameter update group") + result = (await self.init_weights_update_group_communicator(obj))[0] + return result.success, result.message + + async def update_weights_from_distributed( + self: AsyncLLM, + obj: UpdateWeightsFromDistributedReqInput, + ) -> tuple[bool, str]: + self.auto_create_handle_loop() + if self.server_args.mapping.attn.has_dp: + raise RuntimeError("dp_size must be 1 for update weights from distributed") + + # This means that weight sync + # cannot run while requests are in progress. + async with self.model_update_lock.writer_lock: + result = (await self.update_weights_from_distributed_communicator(obj))[0] + return result.success, result.message + + async def update_weights_from_tensor( + self: AsyncLLM, + obj: UpdateWeightsFromTensorReqInput, + ) -> tuple[bool, str]: + self.auto_create_handle_loop() + if self.server_args.mapping.attn.has_dp: + raise RuntimeError("dp_size must be 1 for update weights from tensor") + + # This means that weight sync + # cannot run while requests are in progress. + async with self.model_update_lock.writer_lock: + result = (await self.update_weights_from_tensor_communicator(obj))[0] + return result.success, result.message + + async def get_weights_by_name( + self: AsyncLLM, + obj: GetWeightsByNameReqInput, + ): + self.auto_create_handle_loop() + results = await self.get_weights_by_name_communicator(obj) + all_parameters = [r.parameter for r in results] + if not self.server_args.mapping.attn.has_dp: + return all_parameters[0] + else: + return all_parameters + + async def release_memory_occupation( + self: AsyncLLM, + obj: ReleaseMemoryOccupationReqInput, + ) -> ReleaseMemoryOccupationReqOutput: + self.auto_create_handle_loop() + return (await self.release_memory_occupation_communicator(obj))[0] + + async def resume_memory_occupation( + self: AsyncLLM, + obj: ResumeMemoryOccupationReqInput, + ) -> ResumeMemoryOccupationReqOutput: + self.auto_create_handle_loop() + return (await self.resume_memory_occupation_communicator(obj))[0] + + async def is_sleeping(self: AsyncLLM) -> bool: + self.auto_create_handle_loop() + result = (await self.is_sleeping_communicator(IsSleepingReqInput()))[0] + return result.is_sleeping + + async def get_internal_state(self: AsyncLLM) -> list[dict[Any, Any]]: + req = GetInternalStateReq() + responses: list[GetInternalStateReqOutput] = ( + await self.get_internal_state_communicator(req) + ) + # Many DP ranks + return [res.internal_state for res in responses] + + async def set_internal_state( + self: AsyncLLM, obj: SetInternalStateReq + ) -> list[bool]: + responses: list[SetInternalStateReqOutput] = ( + await self.set_internal_state_communicator(obj) + ) + return [res.updated for res in responses] + + async def get_load(self: AsyncLLM) -> list[GetLoadReqOutput]: + req = GetLoadReqInput() + return await self.get_load_communicator(req) diff --git a/python/tokenspeed/runtime/engine/scheduler_utils.py b/python/tokenspeed/runtime/engine/scheduler_utils.py new file mode 100644 index 0000000..11889eb --- /dev/null +++ b/python/tokenspeed/runtime/engine/scheduler_utils.py @@ -0,0 +1,410 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Helper functions for constructing scheduler specs and events.""" + +import os +from collections.abc import Mapping, Sequence +from typing import Any + +import torch +from tokenspeed_scheduler import ( + Cache, + ExecutionEvent, + ForwardEvent, + PagedCacheGroupConfig, + PagedCacheGroupFamily, + PagedCacheRetention, + PrefixCacheAdjunctSpec, + RequestSpec, + SchedulerConfig, +) + +_CACHE_EVENT_TYPES = { + "WriteBackDoneEvent": Cache.WriteBackDoneEvent, + "PrefetchDoneEvent": Cache.PrefetchDoneEvent, +} +# Emitted only by the flat host tier (FlatMemoryExecutor); the radix executors +# never produce it, so radix behavior is unchanged. hasattr-guarded: the flat +# tier requires a flat-built (post-C3) ext anyway, and an older radix ext must +# keep importing this module. +if hasattr(Cache, "LoadBackDoneEvent"): + _CACHE_EVENT_TYPES["LoadBackDoneEvent"] = Cache.LoadBackDoneEvent +_TRUTHY_ENV_VALUES = {"1", "true", "yes", "on"} + +# Pool-spec string -> scheduler enum (pool_to_paged_cache_groups). +_RETENTION_MAP = { + "full_history": PagedCacheRetention.FullHistory, + "sliding_window": PagedCacheRetention.SlidingWindow, +} +_FAMILY_MAP = { + "history": PagedCacheGroupFamily.History, + "state": PagedCacheGroupFamily.State, +} + + +def make_spec(rid: str, tokens: list[int]) -> RequestSpec: + spec = RequestSpec() + spec.request_id = rid + spec.tokens = tokens + return spec + + +def make_config( + num_device_pages: int, + max_scheduled_tokens: int, + max_batch_size: int, + page_size: int, + num_host_pages: int, + disable_l2_cache: bool, + enable_l3_storage: bool, + prefetch_threshold: int, + role: str, + enable_kv_cache_events: bool = False, + decode_input_tokens: int = 1, + overlap_schedule_depth: int = 0, + disable_prefix_cache: bool = False, + enable_mamba: bool = False, + mamba_cache_chunk_size: int = 64, + mamba_pool_total_chunks: int = 0, + enable_mamba_l2: bool = False, + mamba_l2_host_slots: int = 0, + paged_cache_groups: Sequence["PagedCacheGroupConfig"] | None = None, + enable_mixed_prefill_decode: bool = False, + prefix_cache_adjunct: "PrefixCacheAdjunctSpec | None" = None, +) -> SchedulerConfig: + cfg = SchedulerConfig() + cfg.num_device_pages = num_device_pages + cfg.max_scheduled_tokens = max_scheduled_tokens + cfg.max_batch_size = max_batch_size + cfg.block_size = page_size + + cfg.num_host_pages = num_host_pages + cfg.enable_l3_storage = enable_l3_storage + cfg.prefetch_threshold = prefetch_threshold + cfg.enable_kv_cache_events = enable_kv_cache_events + + if role == "prefill": + cfg.role = SchedulerConfig.Role.P + elif role == "decode": + cfg.role = SchedulerConfig.Role.D + else: + cfg.role = SchedulerConfig.Role.Fused + cfg.decode_input_tokens = decode_input_tokens + cfg.overlap_schedule_depth = overlap_schedule_depth + cfg.disable_prefix_cache = disable_prefix_cache + cfg.disable_l2_cache = disable_l2_cache + + cfg.enable_mamba = enable_mamba + cfg.mamba_cache_chunk_size = mamba_cache_chunk_size + cfg.mamba_pool_total_chunks = mamba_pool_total_chunks + cfg.enable_mamba_l2 = enable_mamba_l2 + cfg.mamba_l2_host_slots = mamba_l2_host_slots + cfg.enable_mixed_prefill_decode = enable_mixed_prefill_decode + if paged_cache_groups: + cfg.paged_cache_groups = list(paged_cache_groups) + # Opt-in; unset means paged-cache groups are transport-only. + if prefix_cache_adjunct is not None: + cfg.prefix_cache_adjunct = prefix_cache_adjunct + return cfg + + +def pool_to_paged_cache_groups(pool: Any) -> list: + """Convert a KV pool's paged_cache_group_specs to scheduler configs.""" + specs = pool.paged_cache_group_specs + if not specs: + return [] + counts = pool.paged_cache_group_page_counts + out = [] + for spec in specs: + retention = _RETENTION_MAP.get(spec.retention) + if retention is None: + raise ValueError( + f"pool_to_paged_cache_groups: unsupported retention " + f"{spec.retention!r} for group {spec.group_id!r}" + ) + family = _FAMILY_MAP.get(spec.family) + if family is None: + raise ValueError( + f"pool_to_paged_cache_groups: unsupported family " + f"{spec.family!r} for group {spec.group_id!r}" + ) + kwargs = dict( + group_id=spec.group_id, + rows_per_page=int(spec.rows_per_page), + entry_stride_tokens=int(spec.entry_stride_tokens), + total_pages=int(counts[spec.group_id]), + retention=retention, + family=family, + ) + if spec.retention == "sliding_window": + kwargs["sliding_window_tokens"] = int(spec.sliding_window_tokens) + out.append(PagedCacheGroupConfig(**kwargs)) + return out + + +def pool_to_prefix_cache_adjunct_spec( + required_group_ids: Sequence[str], +) -> "PrefixCacheAdjunctSpec": + """Build a PrefixCacheAdjunctSpec from required group ids.""" + if not required_group_ids: + raise ValueError( + "pool_to_prefix_cache_adjunct_spec: required_group_ids must be non-empty" + ) + spec = PrefixCacheAdjunctSpec() + spec.required_groups = [str(gid) for gid in required_group_ids] + return spec + + +def should_use_overlap_schedule( + *, + disable_overlap_schedule: bool, + disaggregation_mode: str, +) -> bool: + """Return whether the runtime can use the overlapped scheduler loop.""" + + if disable_overlap_schedule: + return False + if disaggregation_mode in ("prefill", "encode"): + # prefill drain + KV send run only on the non-overlap loop; encode has no LM loop. + return False + return True + + +def make_extend_result_event( + request_id: str, tokens: Sequence[int] = () +) -> "ForwardEvent.ExtendResult": + fe = ForwardEvent.ExtendResult() + fe.request_id = request_id + fe.tokens = list(tokens) + return fe + + +def make_finish_event(request_id: str) -> "ForwardEvent.Finish": + fe = ForwardEvent.Finish() + fe.request_id = request_id + return fe + + +def make_abort_event(request_id: str) -> "ForwardEvent.Abort": + """Finish without caching: AbortEvent skips the radix-tree insert and + never enters Draining, so no host-KV writeback (target or draft) is + issued. Used for numerically-corrupted requests whose KV must not be + reused. + """ + fe = ForwardEvent.Abort() + fe.request_id = request_id + return fe + + +def make_update_reserve_tokens_event(request_id: str, new_reserve_num_tokens: int): + fe = ForwardEvent.UpdateReserveNumTokens() + fe.request_id = request_id + fe.reserve_num_tokens_in_next_schedule_event = new_reserve_num_tokens + return fe + + +def advance_forward(scheduler, forward_events: list) -> None: + ec = ExecutionEvent() + for fe in forward_events: + ec.add_event(fe) + scheduler.advance(ec) + + +def cache_event_to_payload(event) -> dict: + kind = type(event).__name__ + if kind not in _CACHE_EVENT_TYPES: + raise ValueError(f"Unsupported cache event type: {kind}") + return { + "kind": kind, + "op_id": int(event.op_id), + "success": bool(event.success), + "request_id": getattr(event, "request_id", ""), + } + + +def cache_event_from_payload(payload: dict): + kind = payload["kind"] + if kind not in _CACHE_EVENT_TYPES: + raise ValueError(f"Unsupported cache event type: {kind}") + event = _CACHE_EVENT_TYPES[kind]() + event.op_id = int(payload["op_id"]) + event.success = bool(payload["success"]) + request_id = payload.get("request_id", "") + if request_id: + event.request_id = request_id + return event + + +def cache_event_key(payload: dict) -> tuple[str, int]: + return payload["kind"], int(payload["op_id"]) + + +def pop_common_cache_event_payloads( + pending_payloads_by_rank: Sequence[Sequence[dict]], +) -> list[dict]: + if not pending_payloads_by_rank: + return [] + + rank_maps = [] + common_keys = None + for payloads in pending_payloads_by_rank: + rank_map = {cache_event_key(payload): payload for payload in payloads} + rank_maps.append(rank_map) + rank_keys = set(rank_map) + common_keys = rank_keys if common_keys is None else common_keys & rank_keys + if not common_keys: + return [] + + ready_payloads = [] + for key in sorted(common_keys, key=lambda item: (item[1], item[0])): + payload = dict(rank_maps[0][key]) + payload["success"] = all(rank_map[key]["success"] for rank_map in rank_maps) + ready_payloads.append(payload) + return ready_payloads + + +def cache_sync_debug_enabled() -> bool: + value = os.getenv("TS_DEBUG_CACHE_SYNC", "") + return value.strip().lower() in _TRUTHY_ENV_VALUES + + +def _block_tables_from_forward_op( + forward_op: Any, + *, + attr: str, + device: "torch.device | str", + num_reqs: int | None, +) -> dict[str, torch.Tensor]: + raw_tables = getattr(forward_op, attr, None) + if raw_tables is None: + return {} + device = torch.device(device) if isinstance(device, str) else device + items = ( + list(raw_tables.items()) + if isinstance(raw_tables, Mapping) + else list(raw_tables) + ) + out: dict[str, torch.Tensor] = {} + for key_obj, table in items: + key = str(key_obj) + rows = list(table) + if num_reqs is not None and len(rows) != num_reqs: + # No exemption for empty row lists: a silently dropped group + # would hand the flat CUDA-graph replay a per-group hole. + raise ValueError( + f"{attr}[{key}] has {len(rows)} rows but forward op reported " + f"num_reqs={num_reqs}" + ) + if not rows: + # Idle/empty op: callers treat the resulting {} as "no tables". + continue + max_pages = max((len(row) for row in rows), default=0) + if max_pages == 0: + out[key] = torch.empty((len(rows), 0), dtype=torch.int32, device=device) + continue + # One flattened Python list -> single tensor construct (holes stay 0, + # ragged tails pad with -1), instead of O(bs) tiny per-row tensors. + flat_values: list[int] = [] + for row in rows: + row_values = list(row) + flat_values.extend(row_values) + flat_values.extend([-1] * (max_pages - len(row_values))) + # pin_memory as a ctor arg: builds the staging tensor pinned in one + # pass instead of tensor(...).pin_memory()'s second host copy. + flat = torch.tensor( + flat_values, + dtype=torch.int32, + device="cpu", + pin_memory=device.type == "cuda", + ).view(len(rows), max_pages) + out[key] = flat.to(device, non_blocking=True) + return out + + +def paged_cache_block_tables_from_forward_op( + forward_op: Any, + device: "torch.device | str", + *, + num_reqs: int | None = None, +) -> dict[str, torch.Tensor]: + return _block_tables_from_forward_op( + forward_op, + attr="paged_cache_block_tables", + device=device, + num_reqs=num_reqs, + ) + + +def flat_block_tables_from_forward_op( + forward_op: Any, + device: "torch.device | str", + *, + num_reqs: int | None = None, +) -> dict[str, torch.Tensor]: + """Bridge the flat per-group block tables to GPU int32 tensors: absolute + page indices, null hole = 0 preserved, ragged-row padding -1. No + base-offset companion -- the flat path never compacts. + """ + return _block_tables_from_forward_op( + forward_op, + attr="flat_block_tables", + device=device, + num_reqs=num_reqs, + ) + + +def paged_cache_block_table_base_offsets_from_forward_op( + forward_op: Any, + device: "torch.device | str", + *, + num_reqs: int | None = None, +) -> tuple[dict[str, torch.Tensor], dict[str, int]]: + """Convert forward op compact-table base offsets to int32 tensors. + + Returns (gpu_offsets_per_group, cpu_max_per_group). The CPU max is captured + before H2D so callers can size graph-replay buffers without a GPU max + D2H + sync. Empty rows yield max=0; missing keys are absent from the max dict. + """ + raw = getattr(forward_op, "paged_cache_block_table_base_offsets", None) + if raw is None: + return {}, {} + device = torch.device(device) if isinstance(device, str) else device + items = list(raw.items()) if isinstance(raw, Mapping) else list(raw) + out: dict[str, torch.Tensor] = {} + max_per_group: dict[str, int] = {} + for key_obj, offsets in items: + key = str(key_obj) + rows = list(offsets) + if num_reqs is not None and rows and len(rows) != num_reqs: + raise ValueError( + f"paged_cache_block_table_base_offsets[{key}] has {len(rows)} " + f"rows but forward op reported num_reqs={num_reqs}" + ) + if not rows: + max_per_group[key] = 0 + continue + max_per_group[key] = int(max(rows)) + cpu = torch.tensor(rows, dtype=torch.int32, device="cpu") + if device.type == "cuda": + out[key] = cpu.pin_memory().to(device, non_blocking=True) + else: + out[key] = cpu.to(device) + return out, max_per_group diff --git a/python/tokenspeed/runtime/entrypoints/engine.py b/python/tokenspeed/runtime/entrypoints/engine.py new file mode 100755 index 0000000..ccdbe3f --- /dev/null +++ b/python/tokenspeed/runtime/entrypoints/engine.py @@ -0,0 +1,588 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright 2023-2024 SGLang Team +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +""" +The entry point of inference server. + +This file implements python APIs for the inference engine. +""" + +# ruff: noqa: E402 +import asyncio +import atexit +import copy +import dataclasses +import multiprocessing as mp +import os +import signal +import threading +from collections.abc import AsyncIterator, Iterator + +import zmq +import zmq.asyncio + +from tokenspeed.runtime.engine.async_llm import AsyncLLM +from tokenspeed.runtime.engine.llm import LLM + + +def _ignore_threading_atexit(*args, **kwargs) -> None: + return None + + +# Fix a bug of Python threading +setattr(threading, "_register_atexit", _ignore_threading_atexit) + +import torch +import uvloop + +from tokenspeed.runtime.engine.data_parallel_controller import ( + run_data_parallel_controller_process, +) +from tokenspeed.runtime.engine.event_loop import run_event_loop +from tokenspeed.runtime.engine.io_struct import ( + GenerateReqInput, + GetWeightsByNameReqInput, + InitWeightsUpdateGroupReqInput, + ReleaseMemoryOccupationReqInput, + ResumeMemoryOccupationReqInput, + RpcReqInput, + RpcReqOutput, + UpdateWeightFromDiskReqInput, + UpdateWeightsFromDistributedReqInput, + UpdateWeightsFromTensorReqInput, +) +from tokenspeed.runtime.entrypoints.engine_base import EngineBase +from tokenspeed.runtime.utils import ( + MultiprocessingSerializer, + configure_logger, + get_colorful_logger, + launch_dummy_health_check_server, + prepare_model_and_tokenizer, + set_prometheus_multiproc_dir, + set_ulimit, +) +from tokenspeed.runtime.utils.env import envs +from tokenspeed.runtime.utils.process import kill_process_tree +from tokenspeed.runtime.utils.server_args import PortArgs, ServerArgs +from tokenspeed.runtime.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter +from tokenspeed.version import __version__ + +logger = get_colorful_logger(__name__) +asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) + + +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 both run in the main process. + 2. Inter-process communication is done through ICP (each process uses a different port) via the ZMQ library. + """ + + def __init__(self, **kwargs): + """ + The arguments of this function is the same as `tokenspeed/runtime/utils/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("server_args=%r", server_args) + + # Launch subprocesses + tokenizer_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.scheduler_info = scheduler_info + + # Sync facade for blocking callers. Owns its own bg event-loop thread; see runtime/engine/llm.py + # for the queue-bridge semantics. + self.llm = LLM(self.tokenizer_manager) + + def generate( + self, + # The input prompt. It can be a single prompt or a batch of prompts. + prompt: list[str] | str | None = None, + sampling_params: list[dict] | dict | None = None, + # The token ids for text; one can either specify text or input_ids. + input_ids: list[list[int]] | list[int] | None = None, + # SGLang-compatible logprob controls; vLLM-compatible requests use + # sampling_params["logprobs"]. + return_logprob: list[bool] | bool | None = None, + logprob_start_len: list[int] | int | None = None, + top_logprobs_num: list[int] | int | None = None, + token_ids_logprob: list[list[int]] | list[int] | None = None, + return_text_in_logprobs: bool = False, + logprob_format: list[str | None] | str | None = None, + custom_logit_processor: list[str] | str | None = None, + return_hidden_states: bool = False, + stream: bool = False, + bootstrap_host: list[str] | str | None = None, + bootstrap_port: list[int] | int | None = None, + bootstrap_room: list[int] | int | None = None, + data_parallel_rank: int | None = None, + ) -> dict | Iterator[dict]: + """ + The arguments of this function match + ``tokenspeed.runtime.engine.io_struct.GenerateReqInput``. + Please refer to ``GenerateReqInput`` for the documentation. + """ + if self.server_args.mapping.has_attn_dp: + 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.mapping.attn.dp_size: + raise ValueError( + f"data_parallel_rank must be less than dp_size: {self.server_args.mapping.attn.dp_size}" + ) + + obj = GenerateReqInput( + text=prompt, + input_ids=input_ids, + sampling_params=sampling_params, + return_logprob=return_logprob, + logprob_start_len=logprob_start_len, + top_logprobs_num=top_logprobs_num, + token_ids_logprob=token_ids_logprob, + return_text_in_logprobs=return_text_in_logprobs, + logprob_format=logprob_format, + 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, + ) + if stream: + return self.llm.generate_stream(obj) + else: + return self.llm.generate(obj) + + async def async_generate( + self, + # The input prompt. It can be a single prompt or a batch of prompts. + prompt: list[str] | str | None = None, + sampling_params: list[dict] | dict | None = None, + # The token ids for text; one can either specify text or input_ids. + input_ids: list[list[int]] | list[int] | None = None, + input_embeds: torch.Tensor = None, + input_multi_ids: list[list[int]] = None, + input_extra_infos: list[dict] = None, + # Same legacy logprob controls as generate(). + return_logprob: list[bool] | bool | None = None, + logprob_start_len: list[int] | int | None = None, + top_logprobs_num: list[int] | int | None = None, + token_ids_logprob: list[list[int]] | list[int] | None = None, + return_text_in_logprobs: bool = False, + logprob_format: list[str | None] | str | None = None, + custom_logit_processor: list[str] | str | None = None, + return_hidden_states: bool = False, + stream: bool = False, + bootstrap_host: list[str] | str | None = None, + bootstrap_port: list[int] | int | None = None, + bootstrap_room: list[int] | int | None = None, + user_rid: list[str] | str | None = None, + ) -> dict | AsyncIterator[dict]: + """ + The arguments of this function match + ``tokenspeed.runtime.engine.io_struct.GenerateReqInput``. + Please refer to ``GenerateReqInput`` for the documentation. + """ + + obj = GenerateReqInput( + text=prompt, + input_ids=input_ids, + input_embeds=input_embeds, + input_multi_ids=input_multi_ids, + input_extra_infos=input_extra_infos, + sampling_params=sampling_params, + return_logprob=return_logprob, + logprob_start_len=logprob_start_len, + top_logprobs_num=top_logprobs_num, + token_ids_logprob=token_ids_logprob, + return_text_in_logprobs=return_text_in_logprobs, + logprob_format=logprob_format, + 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, + user_rid=user_rid, + ) + generator = self.tokenizer_manager.generate_request(obj) + + async def wrapped_output_generator(original_async_gen): + async for item in original_async_gen: + yield item + + await asyncio.sleep(1) + self.tokenizer_manager.abort_request(obj.rid[0]) + + if stream is True: + return wrapped_output_generator(generator) + else: + return await generator.__anext__() + + def shutdown(self): + """Shutdown the engine""" + # Stop the sync-facade event loop before subprocess teardown so any + # in-flight blocking callers see a clean loop close instead of a + # stale-reference error. + if getattr(self, "llm", None) is not None: + self.llm.shutdown() + 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): + return self.llm.run(self.tokenizer_manager.flush_cache()) + + def pause_scheduler(self, mode: str = "abort"): + """Pause generation (e.g. to swap weights). See AsyncLLM.pause_scheduler.""" + return self.llm.run(self.tokenizer_manager.pause_scheduler(mode=mode)) + + def resume_scheduler(self): + """Resume generation after :meth:`pause_scheduler`.""" + return self.llm.run(self.tokenizer_manager.resume_scheduler()) + + def is_scheduler_paused(self): + """Return whether the scheduler is currently paused.""" + return self.llm.run(self.tokenizer_manager.is_scheduler_paused()) + + def start_profile(self): + self.llm.run(self.tokenizer_manager.start_profile()) + + def stop_profile(self): + self.llm.run(self.tokenizer_manager.stop_profile()) + + def start_expert_distribution_record(self): + self.llm.run(self.tokenizer_manager.start_expert_distribution_record()) + + def stop_expert_distribution_record(self): + self.llm.run(self.tokenizer_manager.stop_expert_distribution_record()) + + def dump_expert_distribution_record(self): + self.llm.run(self.tokenizer_manager.dump_expert_distribution_record()) + + def get_server_info(self): + internal_states = self.llm.run(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, + ) + return self.llm.run(self.tokenizer_manager.init_weights_update_group(obj)) + + 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, + ) + return self.llm.run(self.tokenizer_manager.update_weights_from_distributed(obj)) + + def update_weights_from_tensor( + self, + named_tensors: list[tuple[str, torch.Tensor]], + load_format: str | None = 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.""" + obj = UpdateWeightsFromTensorReqInput( + serialized_named_tensors=[ + MultiprocessingSerializer.serialize(named_tensors) + for _ in range(self.server_args.mapping.world_size) + ], + load_format=load_format, + flush_cache=flush_cache, + ) + return self.llm.run(self.tokenizer_manager.update_weights_from_tensor(obj)) + + def update_weights_from_disk( + self, + model_path: str, + load_format: str | None = 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, + ) + + return self.llm.run(self.tokenizer_manager.update_weights_from_disk(obj)) + + 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) + return self.llm.run(self.tokenizer_manager.get_weights_by_name(obj)) + + def release_memory_occupation(self, tags: list[str] | None = None): + obj = ReleaseMemoryOccupationReqInput(tags=tags) + return self.llm.run(self.tokenizer_manager.release_memory_occupation(obj)) + + def resume_memory_occupation(self, tags: list[str] | None = None): + obj = ResumeMemoryOccupationReqInput(tags=tags) + return self.llm.run(self.tokenizer_manager.resume_memory_occupation(obj)) + + def is_sleeping(self) -> bool: + """Return whether any GPU memory is currently released (data-plane sleep).""" + return self.llm.run(self.tokenizer_manager.is_sleeping()) + + """ + 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) + if not isinstance(recv_req, RpcReqOutput): + raise TypeError(f"Expected RpcReqOutput, got {type(recv_req).__name__}.") + if not recv_req.success: + raise RuntimeError(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 _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" + if not server_args.disable_tf32: + # Force TF32 on for cuBLAS/cuDNN matmuls. setdefault so a user's + # explicit env wins; --disable-tf32 is the documented opt-out. + os.environ.setdefault("NVIDIA_TF32_OVERRIDE", "1") + os.environ.setdefault("TORCH_ALLOW_TF32_CUBLAS_OVERRIDE", "1") + + # Set prometheus env vars + if server_args.enable_metrics: + set_prometheus_multiproc_dir() + + # Set ulimit + set_ulimit() + + # Install a launch-phase SIGQUIT handler so a failing child tears down the + # whole local process tree instead of leaving orphaned workers behind. + # TokenizerManager may replace this handler later during steady-state + # serving. + 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 _launch_subprocesses( + server_args: ServerArgs, port_args: PortArgs | None = None +) -> tuple[AsyncLLM, None, 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) + _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("server_args=%r", server_args) + + # If using model from www.modelscope.cn, first download the model. + server_args.model, server_args.tokenizer = prepare_model_and_tokenizer( + server_args.model, server_args.tokenizer + ) + + scheduler_procs = [] + if not server_args.mapping.attn.has_dp: + # Launch tensor parallel scheduler processes + memory_saver_adapter = TorchMemorySaverAdapter.create( + enable=server_args.enable_memory_saver + ) + + scheduler_pipe_readers = [] + rank_start = server_args.mapping.nprocs_per_node * server_args.node_rank + rank_end = rank_start + server_args.mapping.nprocs_per_node + for rank in range(rank_start, rank_end): + # Create per-rank server_args with rank-initialized mapping + rank_server_args = copy.copy(server_args) + rank_server_args.mapping = copy.deepcopy(server_args.mapping) + rank_server_args.mapping.rank = rank + + reader, writer = mp.Pipe(duplex=False) + + proc = mp.Process( + target=run_event_loop, + args=( + rank_server_args, + port_args, + writer, + ), + ) + 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() + if data.get("status") != "ready": + raise RuntimeError( + "Initialization failed. Please see the error messages above." + ) + + if not envs.TOKENSPEED_BLOCK_NONZERO_RANK_CHILDREN.get(): + # 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( + "Scheduler or DataParallelController %s terminated with %s", + proc.pid, + proc.exitcode, + ) + return None, None, None + + # Launch the main-process async frontend. The detokenizer runs + # inline inside AsyncLLM — no separate subprocess. + tokenizer_manager = AsyncLLM(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( + "Rank %s scheduler is dead. Please check if there are relevant logs.", i + ) + scheduler_procs[i].join() + logger.error("Exit code: %s", 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, None, scheduler_info diff --git a/python/tokenspeed/runtime/entrypoints/engine_base.py b/python/tokenspeed/runtime/entrypoints/engine_base.py new file mode 100755 index 0000000..eba04f3 --- /dev/null +++ b/python/tokenspeed/runtime/entrypoints/engine_base.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the FluentLLM project +# SPDX-FileCopyrightText: Copyright 2023-2024 SGLang Team +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Abstract engine interface for runtime entrypoints.""" + +from abc import ABC, abstractmethod +from collections.abc import Iterator + +import torch + + +class EngineBase(ABC): + """Abstract base class for engine interfaces. + + This interface covers generation, weight updates, and memory control for + both HTTP-based adapters and in-process engines. + """ + + @abstractmethod + def generate( + self, + prompt: list[str] | str | None = None, + sampling_params: list[dict] | dict | None = None, + input_ids: list[list[int]] | list[int] | None = None, + return_logprob: list[bool] | bool | None = None, + logprob_start_len: list[int] | int | None = None, + top_logprobs_num: list[int] | int | None = None, + token_ids_logprob: list[list[int]] | list[int] | None = None, + return_text_in_logprobs: bool = False, + logprob_format: list[str | None] | str | None = None, + custom_logit_processor: list[str] | str | None = None, + return_hidden_states: bool | None = None, + stream: bool | None = None, + bootstrap_host: list[str] | str | None = None, + bootstrap_port: list[int] | int | None = None, + bootstrap_room: list[int] | int | None = None, + data_parallel_rank: int | None = None, + ) -> dict | Iterator[dict]: + """Generate outputs based on given inputs.""" + + @abstractmethod + def flush_cache(self) -> None: + """Flush the cache of the engine.""" + + @abstractmethod + def update_weights_from_tensor( + self, + named_tensors: list[tuple[str, torch.Tensor]], + load_format: str | None = None, + flush_cache: bool = True, + ) -> None: + """Update model weights with in-memory tensor data.""" + + @abstractmethod + def release_memory_occupation(self, tags: list[str] | None = None) -> None: + """Release GPU memory occupation temporarily (optionally by tag).""" + + @abstractmethod + def resume_memory_occupation(self, tags: list[str] | None = None) -> None: + """Resume GPU memory occupation previously released (optionally by tag).""" + + @abstractmethod + def is_sleeping(self) -> bool: + """Return whether any GPU memory is currently released.""" + + @abstractmethod + def shutdown(self) -> None: + """Shutdown the engine and clean up resources.""" diff --git a/python/tokenspeed/runtime/entrypoints/http_server.py b/python/tokenspeed/runtime/entrypoints/http_server.py new file mode 100644 index 0000000..210b073 --- /dev/null +++ b/python/tokenspeed/runtime/entrypoints/http_server.py @@ -0,0 +1,294 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""HTTP server sidecar that runs alongside the smg gateway. + +Runs automatically on ``main_port + 1`` when ``tokenspeed serve`` starts. +Override the port with ``--control-port PORT``. + +Architecture:: + + Client ──► http_server :8001 + ├─ /health, /get_server_info, /get_model_info, + │ /health_check, /abort ──► gRPC engine (direct) + └─ /generate, /v1/*, /flush_cache + ──► smg gateway :8000 ──► gRPC engine +""" + +from __future__ import annotations + +import aiohttp +import grpc +import grpc.aio +import uvicorn +from fastapi import FastAPI, Request +from fastapi.responses import JSONResponse, Response, StreamingResponse +from google.protobuf.json_format import MessageToDict +from smg_grpc_proto.generated import tokenspeed_scheduler_pb2 as pb +from smg_grpc_proto.generated import tokenspeed_scheduler_pb2_grpc as pb_grpc + +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + +app = FastAPI() + +# Set by start() before uvicorn.run(). +_gateway_url: str = "" +_engine_grpc_addr: str = "" +_grpc_channel: grpc.aio.Channel | None = None +_grpc_stub: pb_grpc.TokenSpeedSchedulerStub | None = None + +_STREAM_CHUNK_SIZE = 8192 + +# Proxy timeout: no `total` cap — streaming responses run as long as the model +# keeps generating, and a wall-clock total would abort long generations +# mid-stream. `sock_connect` bounds connecting to smg; `sock_read` bounds +# inactivity (a genuinely hung/stalled upstream) without killing a legitimately +# long but actively-streaming request. +_PROXY_TIMEOUT = aiohttp.ClientTimeout(total=None, sock_connect=30, sock_read=600) + + +def _stub() -> pb_grpc.TokenSpeedSchedulerStub: + # Lazily create a single shared channel/stub. gRPC channels are expensive + # (HTTP/2 connection + background threads) and must be reused, not created + # per request, to avoid leaking sockets and file descriptors. + global _grpc_channel, _grpc_stub + if _grpc_stub is None: + _grpc_channel = grpc.aio.insecure_channel(_engine_grpc_addr) + _grpc_stub = pb_grpc.TokenSpeedSchedulerStub(_grpc_channel) + return _grpc_stub + + +# --------------------------------------------------------------------------- +# Health (proxied to smg → engine) +# --------------------------------------------------------------------------- + + +@app.get("/health") +async def health(request: Request): + return await _proxy_request(request) + + +# --------------------------------------------------------------------------- +# gRPC direct +# --------------------------------------------------------------------------- + + +async def _grpc_call(coro) -> JSONResponse: + """Await a gRPC unary call and serialize the response, mapping engine + errors to a clean 503 instead of an unhandled 500 + stack trace.""" + try: + resp = await coro + except grpc.aio.AioRpcError as exc: + return JSONResponse( + {"error": "engine unavailable", "detail": exc.details()}, + status_code=503, + ) + return JSONResponse(MessageToDict(resp, preserving_proto_field_name=True)) + + +@app.get("/get_server_info") +async def get_server_info(): + return await _grpc_call(_stub().GetServerInfo(pb.GetServerInfoRequest())) + + +@app.get("/get_model_info") +async def get_model_info(): + return await _grpc_call(_stub().GetModelInfo(pb.GetModelInfoRequest())) + + +@app.get("/health_check") +async def health_check(): + return await _grpc_call(_stub().HealthCheck(pb.HealthCheckRequest())) + + +@app.post("/abort") +async def abort(request: Request): + body = await request.json() + return await _grpc_call( + _stub().Abort( + pb.AbortRequest( + request_id=body.get("request_id", ""), + reason=body.get("reason", ""), + ) + ) + ) + + +# --------------------------------------------------------------------------- +# smg passthrough — generation + cache +# --------------------------------------------------------------------------- + + +async def _proxy_request(request: Request) -> StreamingResponse | Response: + url = f"{_gateway_url.rstrip('/')}{request.url.path}" + if request.url.query: + url = f"{url}?{request.url.query}" + body = await request.body() + headers = { + k: v + for k, v in request.headers.items() + if k.lower() not in ("host", "content-length") + } + + # NOTE: the session must outlive the response. For streaming, FastAPI + # consumes the body iterator *after* this function returns, so we cannot + # close the session in an `async with` block here — it would close the + # upstream connection mid-stream. Instead the session is closed in the + # generator's `finally` (streaming) or after `read()` (non-streaming). + session = aiohttp.ClientSession() + try: + resp = await session.request( + method=request.method, + url=url, + headers=headers, + data=body, + timeout=_PROXY_TIMEOUT, + ) + except Exception: + await session.close() + raise + + content_type = resp.headers.get("content-type", "") + if "text/event-stream" in content_type: + + async def _iter(): + try: + async for chunk in resp.content.iter_chunked(_STREAM_CHUNK_SIZE): + yield chunk + finally: + resp.release() + await session.close() + + return StreamingResponse( + _iter(), + status_code=resp.status, + media_type="text/event-stream", + ) + + try: + data = await resp.read() + return Response( + content=data, + status_code=resp.status, + media_type=content_type or "application/json", + ) + finally: + await session.close() + + +@app.api_route("/generate", methods=["GET", "POST"]) +async def generate(request: Request): + return await _proxy_request(request) + + +@app.api_route("/v1/completions", methods=["POST"]) +async def completions(request: Request): + return await _proxy_request(request) + + +@app.api_route("/v1/chat/completions", methods=["POST"]) +async def chat_completions(request: Request): + return await _proxy_request(request) + + +@app.api_route("/v1/models", methods=["GET"]) +async def models(request: Request): + return await _proxy_request(request) + + +@app.api_route("/v1/messages", methods=["POST"]) +async def messages(request: Request): + return await _proxy_request(request) + + +@app.api_route("/v1/responses", methods=["POST"]) +async def responses(request: Request): + return await _proxy_request(request) + + +@app.post("/flush_cache") +async def flush_cache(request: Request): + return await _proxy_request(request) + + +@app.api_route("/start_profile", methods=["GET", "POST"]) +async def start_profile(request: Request): + return await _proxy_request(request) + + +@app.api_route("/stop_profile", methods=["GET", "POST"]) +async def stop_profile(request: Request): + return await _proxy_request(request) + + +# --------------------------------------------------------------------------- +# Server lifecycle +# --------------------------------------------------------------------------- + + +def build_server( + *, + gateway_url: str, + engine_grpc_addr: str, + host: str = "127.0.0.1", + port: int = 8001, +) -> uvicorn.Server: + """Configure the proxy targets and return an unstarted ``uvicorn.Server``. + + The caller runs ``server.run()`` (blocking) and may poll ``server.started`` + to detect when the socket is bound and accepting connections. + + Args: + gateway_url: Base URL of the smg gateway for generation passthrough. + engine_grpc_addr: ``host:port`` of the gRPC engine for direct calls. + host: Bind address. + port: Bind port. + """ + global _gateway_url, _engine_grpc_addr + _gateway_url = gateway_url + _engine_grpc_addr = engine_grpc_addr + logger.info( + "Starting TokenSpeed HTTP server on %s:%d " "(gateway: %s, engine gRPC: %s)", + host, + port, + gateway_url, + engine_grpc_addr, + ) + return uvicorn.Server( + uvicorn.Config(app, host=host, port=port, log_level="warning") + ) + + +def start( + *, + gateway_url: str, + engine_grpc_addr: str, + host: str = "127.0.0.1", + port: int = 8001, +) -> None: + """Start the HTTP server (blocking).""" + build_server( + gateway_url=gateway_url, + engine_grpc_addr=engine_grpc_addr, + host=host, + port=port, + ).run() diff --git a/python/tokenspeed/runtime/execution/breakable_cuda_graph.py b/python/tokenspeed/runtime/execution/breakable_cuda_graph.py new file mode 100644 index 0000000..4bf2cd9 --- /dev/null +++ b/python/tokenspeed/runtime/execution/breakable_cuda_graph.py @@ -0,0 +1,452 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Breakable CUDA graph capture for variable-shape (prefill / extend) forwards. + +A *breakable* CUDA graph captures a forward as an ordered list of zero-arg +callables -- each is either a captured ``CUDAGraph.replay`` (a "graph segment") +or an eager Python function (a "break"). At designated break points (attention / +KV-cache ops, whose metadata is data-dependent and cannot be captured) the +current stream capture is ended, the op runs eagerly, and a fresh segment begins +capturing the remainder. Replay simply calls each segment in order. + +This is the ``torch.compile``-free alternative to piecewise CUDA graphs. The +design is gratefully adapted from vLLM and SGLang, who pioneered the breakable +prefill graph: vLLM's ``BreakableCUDAGraphWrapper`` (the homogeneous segment-list +structure + the ``set_forward_context``/``get_forward_context`` ambient pattern we +mirror in :func:`active_forward`/:func:`current_forward_ctx`) and SGLang's +breakable prefill graph (the eager-copy output handoff at each break). Unlike a +full prefill graph, attention -- the only batch/length-aware op and the source of +the host-side ``max_seq_len_q`` scalar -- stays eager, so it never enters a graph. +Keeping all KV-cache reads/writes in the eager breaks also makes them honor the +per-layer transfer consumer index naturally. + +Address-stability contract (the load-bearing invariant): + +* All segments share one CUDA mempool, so graph-allocated intermediates keep + stable device addresses across replays. +* The runner must copy live inputs into the *same* static input buffers used at + capture before calling :meth:`BreakableCapture.replay`. +* Break-point outputs must land at the *same* address each replay. We achieve + this by allocating a destination buffer in the captured segment (pool-pinned) + and copying the eager op's result into it; the next segment reads that address. +""" + +from __future__ import annotations + +import functools +import gc +from collections import deque +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from typing import Any + +import torch +from tokenspeed_kernel.ops.transform.weak_ref import weak_ref_tensor as _kernel_weak_ref + +__all__ = [ + "BreakableCapture", + "active_forward", + "break_here", + "break_point", + "current_forward_ctx", + "is_breakable_capture_active", + "scrub_padding_tail", + "slice_to_real_tokens", + "weak_ref_tensor", +] + + +# Ambient per-forward ctx; plain module state (one launch thread per rank). +_ambient_ctx: Any = None + + +@contextmanager +def active_forward(ctx: Any) -> Iterator[None]: + """Publish ``ctx`` as the ambient forward context for the enclosed block. + + An eager break runs once at capture and again on every replay, and the args + it closed over at capture are the *dummy* batch's, hence stale. Rather than + thread the live context through ``replay()`` (which would conflate graph + mechanics with forward semantics), the runner wraps capture and each replay + in this, and breaks rebind their captured context arg to the ambient one by + identity (see :func:`break_here`) -- so break bodies read live ``ctx`` + fields exactly like the eager path. Re-entrant (saves/restores the + previous value). + """ + global _ambient_ctx + prev = _ambient_ctx + _ambient_ctx = ctx + try: + yield + finally: + _ambient_ctx = prev + + +def current_forward_ctx() -> Any: + """The ambient forward context, or ``None`` outside an :func:`active_forward`.""" + return _ambient_ctx + + +def weak_ref_tensor(t: Any) -> Any: + """Reference a break-point tensor without pinning its cudagraph mempool slot. + + CUDA tensors are wrapped in a non-owning view (``tokenspeed_kernel`` + ``ops.transform.weak_ref``, an ``at::from_blob`` alias -- the vLLM/sglang + approach), so break closures do not pin pool blocks and graph capture + memory stays ~peak-live instead of scaling with the bucket sum. Safe + because replay is stream-ordered: the captured segment rewrites the + aliased address before the break reads it. Non-tensors and CPU tensors + pass through; if the kernel extension is unavailable this degrades to + the identity (strong ref -- correct, more memory). + """ + if isinstance(t, torch.Tensor) and t.is_cuda: + return _kernel_weak_ref(t) + return t + + +class BreakableCapture: + """Context manager that captures a breakable graph. + + Usage:: + + cap = BreakableCapture(pool=shared_pool) + with cap: + model_forward(...) # attention calls hit break_here() + # later, after copying live inputs into the static buffers: + cap.replay() + + Args: + pool: An optional CUDA mempool id (as returned by + ``torch.cuda.graph_pool_handle()`` or ``CUDAGraph.pool()``) shared by + all segments. If ``None``, the first segment allocates a fresh pool + and the rest reuse it. + stream: An optional dedicated capture stream. CUDA forbids stream capture + on the default stream; if ``None``, a single class-level side stream is + (lazily) created and SHARED by all captures. Sharing one capture stream + is load-bearing for memory: the caching allocator's pool blocks are + stream-keyed, so captures on different streams can never reuse each + other's freed blocks -- a fresh stream per capture makes graph-pool + memory grow with the SUM of bucket sizes instead of the max (measured: + 2478MB -> 564MB for buckets [8192,4096,2048,1024] on the repro). This + mirrors ``torch.cuda.graph``'s shared ``default_capture_stream`` and + its documented "pass the same stream for effective memory sharing". + """ + + _active: BreakableCapture | None = None + _default_capture_stream: torch.cuda.Stream | None = None + + def __init__( + self, pool: Any | None = None, stream: torch.cuda.Stream | None = None + ) -> None: + self.pool = pool + self.segments: list[Callable[[], Any]] = [] + self._current_graph: torch.cuda.CUDAGraph | None = None + self._capturing = False + if stream is None: + if BreakableCapture._default_capture_stream is None: + BreakableCapture._default_capture_stream = torch.cuda.Stream() + stream = BreakableCapture._default_capture_stream + self._stream = stream + self._stream_ctx: Any | None = None + # Break-output handoff buffers keyed by (shape, dtype, device); see break_point. + self._handoff: dict[Any, torch.Tensor] = {} + + @classmethod + def current(cls) -> BreakableCapture | None: + return cls._active + + # -- capture lifecycle ------------------------------------------------- + + def __enter__(self) -> BreakableCapture: + if BreakableCapture.current() is not None: + raise RuntimeError("Nested BreakableCapture is not supported.") + # A GC run during capture invalidates it: destructors of collected + # CUDA graphs call reset, which is illegal while a stream is capturing. + # Clear pending garbage, then keep automatic GC off for the whole + # capture window (restored in __exit__). + gc.collect() + self._gc_was_enabled = gc.isenabled() + gc.disable() + # The capture stream must observe prior entry-stream work (warmup, buffers). + self._stream.wait_stream(torch.cuda.current_stream()) + self._stream_ctx = torch.cuda.stream(self._stream) + self._stream_ctx.__enter__() + BreakableCapture._active = self + self._begin_segment() + return self + + def __exit__(self, *exc: object) -> bool: + try: + self._end_segment() + finally: + BreakableCapture._active = None + if self._stream_ctx is not None: + self._stream_ctx.__exit__(*exc) + self._stream_ctx = None + # Eager breaks ran on the side stream; entry stream must observe them. + torch.cuda.current_stream().wait_stream(self._stream) + if self._gc_was_enabled: + gc.enable() + return False + + def _begin_segment(self) -> None: + assert not self._capturing + graph = torch.cuda.CUDAGraph() + if self.pool is not None: + graph.capture_begin(pool=self.pool) + else: + graph.capture_begin() + self._current_graph = graph + self._capturing = True + + def _end_segment(self) -> None: + if not self._capturing: + return + assert self._current_graph is not None + self._current_graph.capture_end() + self.segments.append(self._current_graph.replay) + # All segments share one pool so intermediate addresses stay stable. + if self.pool is None: + self.pool = self._current_graph.pool() + self._current_graph = None + self._capturing = False + + def add_eager(self, fn: Callable[[], Any]) -> Any: + """End the current segment, run ``fn`` eagerly, record it, start a new one. + + ``fn`` is a zero-arg callable that performs the break-point op and writes + its result into a stable (pool-pinned) address. It is stored verbatim and + re-invoked on every :meth:`replay`. + """ + assert self._capturing, "add_eager called outside an active capture" + self._end_segment() + result = fn() + self.segments.append(fn) + self._begin_segment() + return result + + # -- replay ------------------------------------------------------------ + + def replay(self) -> None: + """Replay all segments in order. + + Breaks read the live forward context from the ambient :func:`active_forward` + scope (the runner wraps replay in it), so this stays a pure graph primitive. + """ + deque((run() for run in self.segments), maxlen=0) + + @property + def num_segments(self) -> int: + return len(self.segments) + + +def is_breakable_capture_active() -> bool: + """True while a :class:`BreakableCapture` is open AND currently capturing.""" + cap = BreakableCapture.current() + return cap is not None and cap._capturing + + +def _record_break( + cap: BreakableCapture, + fn: Callable[..., torch.Tensor], + resolve_dst: Callable[[torch.Tensor], torch.Tensor], + args: tuple, + kwargs: dict, +) -> torch.Tensor: + """Record ``fn(*args, **kwargs)`` as an eager break on ``cap`` (the one closure + builder shared by :func:`break_here` and :func:`break_point`). + + Args/kwargs are bound once at capture time, with two live exceptions: (1) tensor + args alias persistent storage (the static input buffers / pool-pinned segment + intermediates), so they carry live values at replay -- ``weak_ref_tensor`` is + the (currently identity) hook to avoid pinning their pool slots; (2) the + per-forward ``ForwardContext`` is rebound by identity to the live ambient + context each replay (see :func:`active_forward`), so ``fn`` may read live + ``ctx`` fields exactly like the eager path. **Other (loose) non-tensor scalars + are frozen** to their capture-time value -- route per-request quantities + through ``ctx`` / ``forward_*_metadata`` rather than a bare scalar arg. + + ``resolve_dst(result)`` maps the break's first output to its stable handoff + buffer; it is called once (on the capture-time invocation) and the buffer is + reused verbatim on every replay, where the (possibly shorter, see + :func:`_land_in`) live result is copied into it. + """ + weak_args = tuple(weak_ref_tensor(a) for a in args) + weak_kwargs = {k: weak_ref_tensor(v) for k, v in kwargs.items()} + # Capture-time ambient ctx (the dummy batch's), rebound live at replay. + captured_ctx = current_forward_ctx() + state: dict[str, torch.Tensor] = {} + + def replay_fn() -> torch.Tensor: + live_ctx = current_forward_ctx() + + def sub(a: Any) -> Any: + return live_ctx if a is captured_ctx else a + + result = fn( + *(sub(a) for a in weak_args), + **{k: sub(v) for k, v in weak_kwargs.items()}, + ) + dst = state.get("dst") + if dst is None: + dst = state["dst"] = resolve_dst(result) + _land_in(dst, result) + return dst + + return cap.add_eager(replay_fn) + + +def break_here( + fn: Callable[..., torch.Tensor], + dst: torch.Tensor, + *args: Any, + **kwargs: Any, +) -> torch.Tensor: + """Run ``fn(*args, **kwargs)`` as an eager break, landing its result in ``dst``. + + The low-level explicit-destination primitive underneath :func:`break_point` + (which is the decorator every model actually uses -- prefer it; this exists + for callers that must control the handoff buffer's placement themselves, e.g. + a pool-pinned ``dst`` allocated in the current captured segment, and for + exercising the break mechanics directly in unit tests). + + ``dst`` must have a replay-stable address (pool-pinned or persistently owned). + At capture and on every replay, ``fn`` runs eagerly and its result is copied + into ``dst`` (unless ``fn`` already wrote ``dst`` in place and returned it); + the following graph segment reads ``dst``. Outside an active capture this is + a transparent pass-through. Argument binding/freezing semantics are those of + :func:`_record_break`. + + Returns: + ``dst`` (the stable handoff buffer). + """ + cap = BreakableCapture.current() + if cap is None or not cap._capturing: + _land_in(dst, fn(*args, **kwargs)) + return dst + weak_dst = weak_ref_tensor(dst) + return _record_break(cap, fn, lambda _result: weak_dst, args, kwargs) + + +def break_point(method: Callable | None = None) -> Callable: + """Mark a sequence-mixing method as an eager breakable-graph break point. + + Decorate a sequence-mixing method (attention / MLA / linear-mixer / sparse + indexer ``forward``) and it runs as an eager break under a breakable capture -- + the surrounding token-shaped compute (norms, MoE, projections, collectives) is + captured around it automatically, while everything inside the method stays + eager -- or a zero-overhead direct call when not capturing. This is the one + decorator every model uses to mark a break. Use it bare: ``@break_point``. + + The handoff buffer's shape/dtype/device are **inferred from the method's actual + output** at capture time (the break runs during capture regardless), so no + output spec is needed -- it works uniformly for breaks whose output matches no + input (MLA: ``[tokens, heads*v_head_dim]`` vs ``q``'s ``[tokens, heads*qk_head_dim]``) + and for one wrapper that returns different shapes per call (e.g. hybrid full-attn + q-shaped vs GDN z-shaped). Buffers live in a per-capture shape-keyed cache + (:attr:`BreakableCapture._handoff`), shared across same-shape breaks. That + sharing relies on break outputs having strictly sequential lifetimes -- break + K's output is consumed by K's following segment before break K+1 runs (true + for transformer topology, where attention output feeds the adjacent + o-proj/residual). A model whose break output is read by a LATER segment must + not share its shape with an intervening break, or replay silently corrupts. + + Inside the method ``ctx`` is live (rebound by identity at replay), so write the + body exactly like the eager path. Loose non-tensor scalar args are frozen to + their capture-time value -- route per-request quantities through ``ctx`` / metadata. + The decorator never skips the method: 0-row / idle batches remain each decorated + model method's own explicit guard, on the eager path and under capture alike. + """ + + def decorator(method: Callable) -> Callable: + @functools.wraps(method) + def wrapper(*args: Any, **kwargs: Any) -> Any: + # Zero-overhead passthrough off the capture path (no 0-row skipping here). + if not is_breakable_capture_active(): + return method(*args, **kwargs) + cap = BreakableCapture.current() + + def resolve_dst(result: torch.Tensor) -> torch.Tensor: + # Handoff buffer inferred from the capture-time output, shape-keyed. + key = (tuple(result.shape), result.dtype, result.device) + dst = cap._handoff.get(key) + if dst is None: + dst = cap._handoff[key] = torch.empty( + result.shape, dtype=result.dtype, device=result.device + ) + return dst + + return _record_break(cap, method, resolve_dst, args, kwargs) + + return wrapper + + return decorator(method) if method is not None else decorator + + +# -- padded-input helpers (the two strategies of the prefill padding contract) -- + + +def scrub_padding_tail(num_real_tokens: int, *tensors: torch.Tensor | None) -> None: + """Zero the padded tail rows ``[num_real_tokens:]`` of token-shaped tensors in place. + + Under a padded-bucket replay, an eager break receives ``bucket`` rows whose + tail holds garbage (it grows across layers and can overflow to NaN through + projections / FP8 quantize). Zeroing suits breaks whose kernel honors the + live cu-seqlens but whose surrounding ops (varlen attention read, recurrent + scan writeback, FP8 quantize) would otherwise touch the garbage rows. Pass + the real token count from the live metadata's CPU mirror (sync-free on the + launch thread); a no-op on unpadded forwards, and ``None`` tensors are + skipped. + """ + for t in tensors: + if t is not None and num_real_tokens < t.shape[0]: + t[num_real_tokens:].zero_() + + +def slice_to_real_tokens(num_real_tokens: int, *tensors: torch.Tensor | None): + """Return ``tensors`` (in order) each sliced to the real leading rows ``[:num_real_tokens]``. + + The slice-strategy counterpart to :func:`scrub_padding_tail`, for coarse breaks whose + kernel asserts the input row count equals the live metadata token count (e.g. DSA + sparse attention). A tensor already the right length (or ``None``) is returned + unchanged. + """ + return tuple( + t[:num_real_tokens] if (t is not None and num_real_tokens < t.shape[0]) else t + for t in tensors + ) + + +def _land_in(dst: torch.Tensor, result: torch.Tensor) -> None: + """Copy ``result`` into ``dst`` at a stable address. + + ``dst`` is the (possibly token-padded) handoff buffer the next graph segment + reads. ``result`` may cover only the real (unpadded) leading rows -- e.g. a + varlen attention kernel writes only ``sum(cu_seqlens_q)`` rows -- so we copy + into the matching leading slice. Padded rows are left as-is (discarded by the + final output slice). No-op when the op already wrote ``dst`` in place. + """ + if result is dst: + return + if result.shape == dst.shape: + dst.copy_(result) + else: + dst.narrow(0, 0, result.shape[0]).copy_(result) diff --git a/python/tokenspeed/runtime/execution/cache_loc_kernel.py b/python/tokenspeed/runtime/execution/cache_loc_kernel.py new file mode 100644 index 0000000..7330824 --- /dev/null +++ b/python/tokenspeed/runtime/execution/cache_loc_kernel.py @@ -0,0 +1,445 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +""" +Triton kernels for computing cache locations and updating page tables. +""" + +import torch +import triton +import triton.language as tl + +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +@triton.jit +def update_req_to_page_kernel( + # Input pointers + req_pool_indices_ptr, # [batch_size] + new_occupied_pages_ptr, # [total_pages] - flattened + new_occupied_pages_num_ptr, # [batch_size] + pages_copy_starts_ptr, # [batch_size] + cumsum_pages_ptr, # [batch_size] - cumulative sum of new_occupied_pages_num + # Output pointer + req_to_page_ptr, # [req_pool_size+1, context_len] + # Scalars + context_len: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """ + Update req_to_page table with new occupied pages. + Each program handles one request in the batch. + """ + req_idx = tl.program_id(0) + + # Load request metadata + req_pool_idx = tl.load(req_pool_indices_ptr + req_idx) + num_pages = tl.load(new_occupied_pages_num_ptr + req_idx) + copy_start = tl.load(pages_copy_starts_ptr + req_idx) + + # Get offset into flattened new_occupied_pages + offset_idx = tl.where(req_idx > 0, req_idx - 1, 0) + pages_offset = tl.load(cumsum_pages_ptr + offset_idx) + pages_offset = tl.where(req_idx > 0, pages_offset, 0) + + # Process pages in blocks + num_blocks = tl.cdiv(num_pages, BLOCK_SIZE) + for block_idx in range(num_blocks): + block_start = block_idx * BLOCK_SIZE + + # Compute page indices within this block + page_offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = page_offsets < num_pages + + # Load new page IDs + page_ptrs = new_occupied_pages_ptr + pages_offset + page_offsets + new_page_ids = tl.load(page_ptrs, mask=mask, other=0) + + # Compute target positions in req_to_page + target_positions = copy_start + page_offsets + + # Store to req_to_page[req_pool_idx, target_positions] + output_ptrs = req_to_page_ptr + req_pool_idx * context_len + target_positions + tl.store(output_ptrs, new_page_ids, mask=mask) + + +def update_req_to_page( + req_to_page: torch.Tensor, + req_pool_indices: torch.Tensor, + new_occupied_pages: torch.Tensor, + new_occupied_pages_num: torch.Tensor, + pages_copy_starts: torch.Tensor, +) -> None: + """ + Update req_to_page table with new occupied pages using Triton kernel. + + Args: + req_to_page: Request to page table [req_pool_size+1, context_len] + req_pool_indices: Request pool indices [batch_size] + new_occupied_pages: New page IDs [total_pages] - flattened + new_occupied_pages_num: Number of new pages per request [batch_size] + pages_copy_starts: Start position in req_to_page for each request [batch_size] + """ + batch_size = req_pool_indices.shape[0] + context_len = req_to_page.shape[1] + + if new_occupied_pages.shape[0] == 0: + return + + # Compute cumulative sum for offset calculation. + cumsum_pages = torch.cumsum(new_occupied_pages_num, dim=0) + + # Launch kernel - one program per request + BLOCK_SIZE = 128 + grid = (batch_size,) + + update_req_to_page_kernel[grid]( + req_pool_indices, + new_occupied_pages, + new_occupied_pages_num, + pages_copy_starts, + cumsum_pages, + req_to_page, + context_len=context_len, + BLOCK_SIZE=BLOCK_SIZE, + ) + + +@triton.jit +def compute_out_cache_loc_kernel( + # Input pointers + req_pool_indices_ptr, # [batch_size] + input_lengths_ptr, # [batch_size] or None for uniform mode + cache_start_ptr, # [batch_size] + req_to_pages_ptr, # [req_pool_size+1, max_pages] + cumsum_lengths_ptr, # [batch_size] or None for uniform mode + # Output pointer + out_cache_loc_ptr, # [total_tokens] + # Scalars + uniform_input_length, # used when input_lengths_ptr is None + page_size: tl.constexpr, + max_pages: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """ + Unified kernel to compute out_cache_loc for both prefill and decode. + + For each token in each request, compute: + position = cache_start[req_idx] + token_offset_in_seq + page_idx = position // page_size + offset_in_page = position % page_size + page_id = req_to_pages[req_pool_idx, page_idx] + out_cache_loc = page_id * page_size + offset_in_page + + For decode, input_lengths are all 1. + For prefill, input_lengths vary. + + When all requests share the same input_length (the multi-step drafter + case), callers pass ``input_lengths_ptr=None`` (and ``cumsum_lengths_ptr=None``) + together with ``uniform_input_length`` set to the shared length. Triton + specializes the kernel on the None-ness of the pointers at JIT time and + dead-code-eliminates the corresponding GMEM reads. + """ + # Program ID represents which request we're processing + req_idx = tl.program_id(0) + + # Load request metadata. + req_pool_idx = tl.load(req_pool_indices_ptr + req_idx) + valid_cache_len = tl.load(cache_start_ptr + req_idx) + + if input_lengths_ptr is not None: + input_length = tl.load(input_lengths_ptr + req_idx) + # Always load from cumsum, use 0 index for first request to ensure type consistency + offset_idx = tl.where(req_idx > 0, req_idx - 1, 0) + output_offset = tl.load(cumsum_lengths_ptr + offset_idx) + # Zero out offset for first request + output_offset = tl.where(req_idx > 0, output_offset, 0) + else: + input_length = uniform_input_length + output_offset = req_idx * uniform_input_length + + # Process tokens in blocks + num_blocks = tl.cdiv(input_length, BLOCK_SIZE) + for block_idx in range(num_blocks): + block_start = block_idx * BLOCK_SIZE + + # Compute token offsets within this block + token_offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = token_offsets < input_length + + # Compute logical positions + positions = valid_cache_len + token_offsets + + # Compute page indices and offsets + page_indices = positions // page_size + overflow = page_indices >= max_pages + # Clamp to last valid page to avoid OOB GMEM read. + page_indices = tl.minimum(page_indices, max_pages - 1) + offsets_in_page = positions % page_size + + # Load page IDs from req_to_pages + # req_to_pages is [req_pool_size+1, max_pages] + page_ptrs = req_to_pages_ptr + req_pool_idx * max_pages + page_indices + page_ids = tl.load(page_ptrs, mask=mask, other=0) + + # Compute physical cache locations + cache_locs = page_ids * page_size + offsets_in_page + # For overflow tokens, route to slot 0 (a fixed safe dummy target that + # never aliases a real request's KV data). This avoids using a dynamic + # req_to_pages[0][0] load whose value can change at runtime and corrupt + # other requests' KV cache or trigger IndexKernel out-of-bounds. + cache_locs = tl.where(overflow, 0, cache_locs) + + # Store to output + output_ptrs = out_cache_loc_ptr + output_offset + token_offsets + tl.store(output_ptrs, cache_locs, mask=mask) + + +def compute_out_cache_loc( + out_cache_loc_ptr, + req_pool_indices: torch.Tensor, # [batch_size] + input_lengths: torch.Tensor, # [batch_size] + cache_start: torch.Tensor, # [batch_size] + req_to_pages: torch.Tensor, # [req_pool_size+1, max_pages] + page_size: int, +) -> None: + batch_size = req_pool_indices.shape[0] + max_pages = req_to_pages.shape[1] + + cumsum_lengths = torch.cumsum(input_lengths, dim=0) + + BLOCK_SIZE = 128 + grid = (batch_size,) + + compute_out_cache_loc_kernel[grid]( + req_pool_indices, + input_lengths, + cache_start, + req_to_pages, + cumsum_lengths, + out_cache_loc_ptr, + 0, # uniform_input_length unused when input_lengths_ptr is not None + page_size=page_size, + max_pages=max_pages, + BLOCK_SIZE=BLOCK_SIZE, + ) + + +@triton.jit +def fused_decode_input_prep_kernel( + # Inputs + req_pool_indices_ptr, # [batch_size] + valid_cache_lengths_ptr, # [req_pool_size+1] + req_to_pages_ptr, # [req_pool_size+1, max_pages] + # Outputs + out_cache_loc_ptr, # [batch_size * uniform_input_length] + positions_ptr, # [batch_size * uniform_input_length] + seq_lens_out_ptr, # [batch_size] + # Scalars + uniform_input_length, + page_size: tl.constexpr, + max_pages: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + """One launch fuses the decode-uniform path's four small kernels. + + Replaces: + valid_cache_lengths.index_select(0, req_pool_indices) + compute_out_cache_loc_uniform + compute_position_triton (decode branch) + torch.add(input_lengths, valid_cache_lengths, out=seq_lens) + + Each program handles one request. We do one GMEM read of + `valid_cache_lengths[pool_idx]` and reuse it for the seq_lens write, + the position writes, and the out_cache_loc page-table lookup. + """ + req_idx = tl.program_id(0) + pool_idx = tl.load(req_pool_indices_ptr + req_idx) + cache_start = tl.load(valid_cache_lengths_ptr + pool_idx) + + # seq_lens[req_idx] = cache_start + uniform_input_length + tl.store(seq_lens_out_ptr + req_idx, cache_start + uniform_input_length) + + output_offset = req_idx * uniform_input_length + + num_blocks = tl.cdiv(uniform_input_length, BLOCK_SIZE) + for block_idx in range(num_blocks): + block_start = block_idx * BLOCK_SIZE + token_offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = token_offsets < uniform_input_length + + positions_local = cache_start + token_offsets + page_indices = positions_local // page_size + overflow = page_indices >= max_pages + # Clamp to last valid page to avoid OOB GMEM read. + page_indices = tl.minimum(page_indices, max_pages - 1) + offsets_in_page = positions_local % page_size + + page_ptrs = req_to_pages_ptr + pool_idx * max_pages + page_indices + page_ids = tl.load(page_ptrs, mask=mask, other=0) + cache_locs = page_ids * page_size + offsets_in_page + # Route overflow tokens to slot 0 (fixed safe dummy target). + cache_locs = tl.where(overflow, 0, cache_locs) + + tl.store( + out_cache_loc_ptr + output_offset + token_offsets, + cache_locs, + mask=mask, + ) + tl.store( + positions_ptr + output_offset + token_offsets, + positions_local, + mask=mask, + ) + + +def fused_decode_input_prep( + out_cache_loc_ptr, + positions_ptr, + seq_lens_out_ptr, + req_pool_indices: torch.Tensor, # [batch_size] + valid_cache_lengths: torch.Tensor, # [req_pool_size+1] + uniform_input_length: int, + req_to_pages: torch.Tensor, # [req_pool_size+1, max_pages] + page_size: int, +) -> None: + """Decode-only fast path: one Triton launch writes out_cache_loc, + positions, and seq_lens, reading `valid_cache_lengths[pool_idx]` + directly so the per-iter indexSelect + add are gone too. + """ + batch_size = req_pool_indices.shape[0] + max_pages = req_to_pages.shape[1] + BLOCK_SIZE = 128 + grid = (batch_size,) + fused_decode_input_prep_kernel[grid]( + req_pool_indices, + valid_cache_lengths, + req_to_pages, + out_cache_loc_ptr, + positions_ptr, + seq_lens_out_ptr, + uniform_input_length, + page_size=page_size, + max_pages=max_pages, + BLOCK_SIZE=BLOCK_SIZE, + ) + + +def compute_out_cache_loc_uniform( + out_cache_loc_ptr, + req_pool_indices: torch.Tensor, # [batch_size] + uniform_input_length: int, + cache_start: torch.Tensor, # [batch_size] + req_to_pages: torch.Tensor, # [req_pool_size+1, max_pages] + page_size: int, +) -> None: + """Specialized entry point when every request has the same ``input_length``. + + Skips the per-call ``torch.full`` + ``cumsum`` host-side work and the + corresponding GMEM reads inside the kernel. Used by the multi-step drafter + where each request decodes exactly ``spec_num_steps - 1`` tokens. + """ + batch_size = req_pool_indices.shape[0] + max_pages = req_to_pages.shape[1] + + BLOCK_SIZE = 128 + grid = (batch_size,) + + compute_out_cache_loc_kernel[grid]( + req_pool_indices, + None, # input_lengths_ptr is None → kernel uses uniform_input_length + cache_start, + req_to_pages, + None, # cumsum_lengths_ptr is None → kernel computes offset analytically + out_cache_loc_ptr, + uniform_input_length, + page_size=page_size, + max_pages=max_pages, + BLOCK_SIZE=BLOCK_SIZE, + ) + + +def update_block_table(forward_op, device, req_to_page): + def flatten_and_to_device(data, dtype=torch.int32): + if not data: + return torch.tensor([], dtype=dtype, device=device) + + # Flatten one level if data is a list of lists + if isinstance(data[0], (list, tuple)): + flat = [x for inner in data for x in inner] + else: + flat = data + + if not flat: + return torch.tensor([], dtype=dtype, device=device) + + tensor = torch.tensor(flat, dtype=dtype, device="cpu", pin_memory=True) + return tensor.to(device, non_blocking=True) + + # sizes[i] is the number of newly allocated pages for request i. + if all(n == 0 for n in forward_op.sizes): + return + + max_pages = req_to_page.shape[1] + # Clamp a request that would overflow req_to_page instead of crashing the + # engine. Happens when MTP accept-rate collapse keeps a request alive past + # context_len; its KV drops but it will be finished shortly. + sizes = list(forward_op.sizes) + begins = list(forward_op.begins) + # new_occupied_pages is a list-of-lists [batch, size_i] of page ids; + # take a shallow copy so we can trim the offending request's row. + new_occupied_pages = [list(row) for row in forward_op.new_occupied_pages] + request_ids = list(forward_op.request_ids) + for i, (begin, size) in enumerate(zip(begins, sizes)): + if begin + size > max_pages: + clamped = max(0, max_pages - begin) + logger.warning( + "page copy would exceed req_to_page capacity for req %s: " + "begin=%s + size=%s = %s > req_to_page.shape[1]=%s; " + "clamping size to %s to avoid engine crash. The request is past " + "its context-length bound and will be finished by the length " + "check; KV writes after this point are dropped.", + request_ids[i] if i < len(request_ids) else "?", + begin, + size, + begin + size, + max_pages, + clamped, + ) + sizes[i] = clamped + # Keep new_occupied_pages[i] consistent with the clamped size so + # the kernel's cumsum-based offsets stay aligned across the batch. + new_occupied_pages[i] = new_occupied_pages[i][:clamped] + + new_occupied_pages_num = flatten_and_to_device(sizes, dtype=torch.int32) + pages_copy_starts = flatten_and_to_device(begins, dtype=torch.int32) + new_occupied_pages_t = flatten_and_to_device(new_occupied_pages, dtype=torch.int32) + request_pool_indices = flatten_and_to_device( + forward_op.request_pool_indices, dtype=torch.int64 + ) + update_req_to_page( + req_to_page=req_to_page, + req_pool_indices=request_pool_indices, + new_occupied_pages=new_occupied_pages_t, + new_occupied_pages_num=new_occupied_pages_num, + pages_copy_starts=pages_copy_starts, + ) diff --git a/python/tokenspeed/runtime/execution/context.py b/python/tokenspeed/runtime/execution/context.py new file mode 100644 index 0000000..743f299 --- /dev/null +++ b/python/tokenspeed/runtime/execution/context.py @@ -0,0 +1,107 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import torch + +from tokenspeed.runtime.execution.forward_batch_info import ( + CaptureHiddenMode, + ForwardMode, +) + +if TYPE_CHECKING: + from tokenspeed.runtime.layers.attention.backends.base import AttentionBackend + from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool + + +@dataclass +class ForwardContext: + """Do not contain Tensor""" + + # --- attention infrastructure --- + attn_backend: AttentionBackend + token_to_kv_pool: BaseTokenToKVPool + + # --- meta data --- + bs: int + num_extends: int + input_num_tokens: int + forward_mode: ForwardMode | None + req_to_page: torch.Tensor | None = None + capture_hidden_mode: CaptureHiddenMode | None = CaptureHiddenMode.NULL + # Normalized explicit decode input overrides for this forward, if any. + decode_input_ids: list[int] | None = None + + # --- dp attention --- + global_num_tokens: list[int] | None = None + global_bs: list[int] | None = None + all_decode_or_idle: bool = False + all_extend: bool = False + # Models that need specific collective sizing (e.g. draft models whose + # first-step forward narrows activations) report these via + # ``report_collective_sizing``. Unset (None) means comm sizing falls + # back to ``input_num_tokens`` / ``global_num_tokens``. + collective_num_tokens: int | None = None + collective_global_num_tokens: list[int] | None = None + + # --- logits processor --- + gather_ids: torch.Tensor | None = None + + # --- spec-decode draft (drafter-owned buffers plumbed per forward) --- + # draft_seq_lens_buf: mutable per-request seq_lens alias the draft backend reads. + draft_seq_lens_buf: torch.Tensor | None = None + # accept_lengths: per-request accepted verify width for cache_seqlens correction. + accept_lengths: torch.Tensor | None = None + + # DSA sparse top-k shared across layers and draft steps. + dsa_prefill_topk: Any | None = None + dsa_decode_topk: Any | None = None + + # DSA SWA slot mapping + compressor memo, computed once per forward, shared across layers. + dsa_swa_slot_mapping: torch.Tensor | None = None + dsa_compressor_slot_cache: Any | None = None + + +@contextmanager +def report_collective_sizing( + ctx: ForwardContext, + num_tokens: int, + global_num_tokens: list[int] | None, +): + """Report model-specific collective sizing for the duration of the scope. + + When a model needs to specify particular collective token counts (e.g. + draft models narrowing activations to one row per request), wrap the + model forward in this context manager. Comm collectives will use the + reported values instead of ``input_num_tokens`` / ``global_num_tokens``. + Automatically cleared on exit so later forwards use the default sizing. + """ + ctx.collective_num_tokens = num_tokens + ctx.collective_global_num_tokens = global_num_tokens + try: + yield + finally: + ctx.collective_num_tokens = None + ctx.collective_global_num_tokens = None diff --git a/python/tokenspeed/runtime/execution/cuda_graph_wrapper.py b/python/tokenspeed/runtime/execution/cuda_graph_wrapper.py new file mode 100644 index 0000000..94c52d3 --- /dev/null +++ b/python/tokenspeed/runtime/execution/cuda_graph_wrapper.py @@ -0,0 +1,1170 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import bisect +import gc +import queue +from collections.abc import Callable +from contextlib import contextmanager +from typing import TYPE_CHECKING + +import torch +import torch.distributed as dist +import tqdm + +from tokenspeed.runtime.configs.paged_cache_spec import ( + compute_max_logical_pages_for_capture, +) +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.execution.forward_batch_info import ( + CaptureHiddenMode, + ForwardMode, +) +from tokenspeed.runtime.layers.attention.backends.base import ( + init_backend_cuda_graph_state, +) +from tokenspeed.runtime.sampling.backends.base import CUDA_GRAPH_VARIANT_DEFAULT +from tokenspeed.runtime.sampling.sampling_batch_info import SamplingBatchInfo +from tokenspeed.runtime.utils import ( + get_available_gpu_memory, + get_colorful_logger, +) +from tokenspeed.runtime.utils.nvtx import nvtx_range + +if TYPE_CHECKING: + from tokenspeed.runtime.execution.drafter.base import BaseDrafter + from tokenspeed.runtime.execution.input_buffer import InputBuffers + from tokenspeed.runtime.execution.model_executor import ModelExecutorConfig + from tokenspeed.runtime.execution.runtime_states import RuntimeStates + from tokenspeed.runtime.layers.attention.backends.base import AttentionBackend + from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool + from tokenspeed.runtime.sampling.backends.base import SamplingBackend + +logger = get_colorful_logger(__name__) + + +_is_capture_mode = False + + +def get_is_capture_mode() -> bool: + return _is_capture_mode + + +def _should_update_mamba_state_after_mtp_verify( + drafter, attn_backend, forward_mode: ForwardMode +) -> bool: + return ( + drafter is not None + and forward_mode.is_decode() + and hasattr(attn_backend, "update_mamba_state_after_mtp_verify") + ) + + +@contextmanager +def freeze_gc(enable_cudagraph_gc: bool): + """ + Optimize garbage collection during CUDA graph capture. + Clean up, then freeze all remaining objects from being included + in future collections if GC is disabled during capture. + """ + gc.collect() + should_freeze = not enable_cudagraph_gc + if should_freeze: + gc.freeze() + try: + yield + finally: + if should_freeze: + gc.unfreeze() + gc.collect() + + +def get_batch_sizes_to_capture(config: ModelExecutorConfig): + capture_bs = config.cudagraph_capture_sizes + max_bs = config.max_num_seqs // max(config.data_parallel_size, 1) + + if capture_bs is None: + if config.disable_cuda_graph_padding: + capture_bs = list(range(1, 33)) + [64, 96, 128, 160] + else: + capture_bs = [1, 2, 4] + [i * 8 for i in range(1, 21)] + + if max(capture_bs) > max_bs: + capture_bs = list(sorted(set(capture_bs + [max_bs - 1] + [max_bs]))) + + effective_max = min(config.max_cudagraph_capture_size, max_bs) + capture_bs = [bs for bs in capture_bs if 0 < bs <= effective_max] + return capture_bs + + +global_graph_memory_pool = None + + +class DeepEPCudaGraphRunnerAdapter: + """Manages DeepEP dispatch mode consistency across CUDA graph capture/replay. + + During capture the forward pass (including DeepEP low-latency RDMA + dispatch/combine) is recorded. On replay the Python wrapper code + that normally sets dispatch mode and manages the RDMA workspace + never re-executes. This adapter restores both before each replay. + + Follows the same CUDA graph replay contract as the upstream DeepEP runner. + """ + + def __init__(self): + self._active = False + + @staticmethod + def _get_buffer_cls(): + try: + from tokenspeed_kernel.ops.communication.deep_ep import ( + DeepEPBuffer, + ) + + return DeepEPBuffer + except ImportError: + return None + + def capture(self): + """Call before ``torch.cuda.graph()`` capture.""" + cls = self._get_buffer_cls() + if cls is None or cls._buffer is None: + return + self._active = True + cls.set_dispatch_mode_as_low_latency() + + def replay(self): + """Call before every ``graph.replay()``; restores dispatch mode + and resets RDMA workspace so stale sync state doesn't corrupt + the combine kernel across replays.""" + if not self._active: + return + cls = self._get_buffer_cls() + if cls is None or cls._buffer is None: + return + cls.set_dispatch_mode_as_low_latency() + cls.clean_buffer() + + +class CudaGraphWrapper: + """ + Wraps a forward_func and transparently dispatches to either a captured + CUDA graph (decode, supported batch size) or the eager path (prefill / + unsupported batch size). + + Callers always use the same interface:: + + output_tokens, output_lengths, output_logprobs = runner( + bs, ctx, sampling_info, req_to_page, + extend_with_prefix=..., extend_prefix_lens=..., + ) + + Internally the wrapper owns both paths and calls init_forward_metadata + with use_cuda_graph=True/False to select the appropriate backend buffers. + """ + + def __init__( + self, + forward_func: Callable, + attn_backend: AttentionBackend, + token_to_kv_pool: BaseTokenToKVPool, + input_buffers: InputBuffers, + config: ModelExecutorConfig, + draft_attn_backend: AttentionBackend | None = None, + draft_token_to_kv_pool: BaseTokenToKVPool | None = None, + drafter: BaseDrafter | None = None, + capturable_grammar=None, + eager_grammar_buffers=None, + sampling_backend: SamplingBackend | None = None, + runtime_states: RuntimeStates | None = None, + ): + self.config = config + self.attn_backend = attn_backend + self.draft_attn_backend = draft_attn_backend + self.draft_token_to_kv_pool = draft_token_to_kv_pool + self.token_to_kv_pool = token_to_kv_pool + self.drafter = drafter + self.sampling_backend = sampling_backend + self.input_buffers = input_buffers + self.capturable_grammar = capturable_grammar + self.eager_grammar_buffers = eager_grammar_buffers + self.runtime_states = runtime_states + self.enable_torch_compile = getattr(config, "enable_torch_compile", False) + self.disable_padding = config.disable_cuda_graph_padding + self.enable_cudagraph_gc = getattr(config, "enable_cudagraph_gc", True) + self.device = config.device + self.gpu_id = config.gpu_id + self.global_rank = config.global_rank + self.context_len = config.context_len + self.vocab_size = config.vocab_size + self.grammar_backend = config.grammar_backend + self.capture_bs = get_batch_sizes_to_capture(config) + self.max_bs = max(self.capture_bs) + self.max_tokens_per_req = ( + config.spec_num_tokens if config.spec_algo is not None else 1 + ) + self.overlap_schedule_depth = config.overlap_schedule_depth + self.use_v4_mtp_paged_metadata = config.use_v4_mtp_paged_metadata + self.dp_size = config.data_parallel_size + self.world_size = config.world_size + # Backends alias their cache_seqlens buffer. Draft backend aliases + # the drafter-owned draft_seq_lens to keep InputBuffers read-only. + init_backend_cuda_graph_state( + attn_backend, + self.max_bs, + self.input_buffers.seq_lens_buf, + paged_cache_group_specs=tuple(token_to_kv_pool.paged_cache_group_specs), + max_tokens_per_req=self.max_tokens_per_req, + overlap_schedule_depth=self.overlap_schedule_depth, + ) + if draft_attn_backend is not None: + init_backend_cuda_graph_state( + draft_attn_backend, + self.max_bs, + self.drafter.draft_seq_lens_buf, + paged_cache_group_specs=tuple( + draft_token_to_kv_pool.paged_cache_group_specs + ), + max_tokens_per_req=self.max_tokens_per_req, + overlap_schedule_depth=self.overlap_schedule_depth, + ) + + # Drafter (Eagle) is constructed with the target's req_to_page + # (ModelExecutor passes the same self.req_to_page to both), and the + # replay path hands both backends the same req_pool_indices. The + # block-table gather is req_to_page[req_pool_indices] (see + # _create_block_kv_indices; it does not depend on seq_lens), so both + # backends would compute identical block_kv_indices. When the backing + # buffer shapes/dtypes also line up, point the draft backend at the + # target's buffer and skip its gather+copy in the replay path: the + # target's metadata prep runs first and populates the shared buffer + # (see init_forward_metadata_replay_cuda_graph). + target_kv = getattr(attn_backend, "decode_cuda_graph_kv_indices", None) + draft_kv = getattr(draft_attn_backend, "decode_cuda_graph_kv_indices", None) + if ( + target_kv is not None + and draft_kv is not None + and target_kv.shape == draft_kv.shape + and target_kv.dtype == draft_kv.dtype + ): + draft_attn_backend.decode_cuda_graph_kv_indices = target_kv + draft_attn_backend._block_table_aliased = True + + self.graph_variants = ( + sampling_backend.cuda_graph_capture_variants(self.max_tokens_per_req) + if sampling_backend is not None + else (CUDA_GRAPH_VARIANT_DEFAULT,) + ) + self.graphs: dict[tuple[str, int], torch.cuda.CUDAGraph] = {} + self.output_buffers: dict[tuple[str, int], tuple] = {} + + self._forward_func: Callable | None = forward_func + self.disable = config.enforce_eager + self.deepep_adapter = DeepEPCudaGraphRunnerAdapter() + if not self.disable: + self.capture() + + # ------------------------------------------------------------------ + # Graph capture + # ------------------------------------------------------------------ + + def capture(self): + """ + Capture CUDA graphs for all configured batch sizes. + + Args: + forward_func: ModelExecutor.forward_step(bs, ctx, sampling_info). + """ + rank = self.global_rank + with freeze_gc(self.enable_cudagraph_gc): + self.stream = torch.cuda.Stream() + # Capture backend-declared sampler variants explicitly. + capture_items = [ + (variant, bs) + for variant in self._cuda_graph_capture_variants() + for bs in self.capture_bs + ] + capture_range = tqdm.tqdm(capture_items) if rank == 0 else capture_items + if rank == 0: + logger.info("Capturing batches: %s", self.capture_bs) + for variant, bs in capture_range: + if rank == 0: + avail_mem = get_available_gpu_memory( + self.device, self.gpu_id, empty_cache=False + ) + variant_desc = ( + "" + if variant == CUDA_GRAPH_VARIANT_DEFAULT + else f" variant={variant}" + ) + capture_range.set_description( + f"Capturing batches ({bs=}{variant_desc} {avail_mem=:.2f} GB)" + ) + graph, output_buffers = self._capture_one(bs, variant=variant) + self.graphs[(variant, bs)] = graph + self.output_buffers[(variant, bs)] = output_buffers + + def _cuda_graph_capture_variants(self) -> tuple[str, ...]: + if self.sampling_backend is None: + return (CUDA_GRAPH_VARIANT_DEFAULT,) + variants = self.sampling_backend.cuda_graph_capture_variants( + self.max_tokens_per_req + ) + if not variants: + return (CUDA_GRAPH_VARIANT_DEFAULT,) + deduped = tuple(dict.fromkeys((CUDA_GRAPH_VARIANT_DEFAULT, *variants))) + return deduped + + def _prepare_sampling_capture(self, bs: int, variant: str) -> None: + if self.sampling_backend is None: + return + self.sampling_backend.prepare_capture_variant( + bs=bs, + num_tokens_per_req=self.max_tokens_per_req, + variant=variant, + ) + + def _cuda_graph_replay_variant(self) -> str: + if self.sampling_backend is None: + return CUDA_GRAPH_VARIANT_DEFAULT + return self.sampling_backend.cuda_graph_replay_variant(self.max_tokens_per_req) + + def _cuda_graph_key(self, bs: int) -> tuple[str, int]: + variant = self._cuda_graph_replay_variant() + key = (variant, bs) + if key in self.graphs: + return key + if variant != CUDA_GRAPH_VARIANT_DEFAULT: + captured_variants = sorted( + graph_variant + for graph_variant, graph_bs in self.graphs + if graph_bs == bs + ) + raise RuntimeError( + "Sampling backend requested CUDA graph variant " + f"{variant!r} for batch size {bs}, but it was not captured. " + f"Captured variants for this batch size: {captured_variants}." + ) + return (CUDA_GRAPH_VARIANT_DEFAULT, bs) + + def _has_cuda_graph_for_bs(self, bs: int) -> bool: + return (CUDA_GRAPH_VARIANT_DEFAULT, bs) in self.graphs + + def _capture_one(self, bs: int, variant: str = CUDA_GRAPH_VARIANT_DEFAULT): + graph = torch.cuda.CUDAGraph() + + capture_forward_mode = ForwardMode.DECODE + ctx = ForwardContext( + attn_backend=self.attn_backend, + token_to_kv_pool=self.token_to_kv_pool, + bs=bs, + num_extends=0, + input_num_tokens=bs * self.max_tokens_per_req, + forward_mode=capture_forward_mode, + capture_hidden_mode=( + CaptureHiddenMode.FULL + if self.drafter is not None + else CaptureHiddenMode.NULL + ), + ) + + # For DP mode, global_num_tokens must be set so that the MoE + # all-gather comm layers know token counts for all DP ranks. + # During capture, use uniform dummy counts across ranks. + if self.dp_size > 1: + ctx.global_num_tokens = [bs * self.max_tokens_per_req] * self.world_size + # global_bs must ALSO be set at capture. The draft first step's + # collective sizing (reported via report_collective_sizing) reads + # global_bs; if left None at capture it records a single-rank + # layout (fallback branch in comm_manager), but at replay global_bs + # is the live per-rank batch list -> multi-rank layout. The mismatch + # makes the captured (frozen-offset) gather read uninitialized + # symm-mem -> NaN draft logits -> accept_rate 0. Set the matching + # uniform dummy. + ctx.global_bs = [bs] * self.world_size + + # Capture with is_all_greedy=False so the graph records the full + # top_k_top_p_sampling path (greedy-only requests are served by the + # same path with top_k=1 in the buffer, which effectively argmaxes). + # is_all_greedy=True at capture would freeze the graph into + # argmax and bypass per-request seeding at replay. + ibd = self.input_buffers + sampling_info = SamplingBatchInfo( + req_pool_indices=ibd.req_pool_indices_buf[:bs], + valid_cache_lengths=( + self.runtime_states.valid_cache_lengths + if self.runtime_states is not None + else None + ), + is_all_greedy=False, + vocab_size=self.vocab_size, + device=self.device, + ) + + from tokenspeed.runtime.grammar.capturable_grammar import ( + bind_grammar_mask_buf, + ) + + # Bind whichever grammar buffer is active so the captured sampler + # records the apply_vocab_mask call. At replay, runtime fills the + # bound buffer in place (hostfunc for capturable, sync H2D for + # eager) — the captured graph reads from the same memory. + bind_grammar_mask_buf( + sampling_info, + self.eager_grammar_buffers, + bs, + spec=self.drafter is not None, + capturable=self.capturable_grammar, + grammar_backend=self.grammar_backend, + ) + + def run_once(): + # Dummy add_batch keeps the grammar queue 1:1 with replays — + # fetch_batch pops once per forward, so warmup + capture + # would otherwise raise queue.Empty. + if self.capturable_grammar is not None: + self.capturable_grammar.add_batch( + grammars=[None] * bs, bs=bs, has_candidates=False + ) + return self._forward_func(bs=bs, ctx=ctx, sampling_info=sampling_info) + + # Warm up before capture. + for _ in range(4): + torch.cuda.synchronize() + dist.barrier() + self._prepare_sampling_capture(bs=bs, variant=variant) + # Keep warmup seq_lens >= q_len_per_req so no query row gets an + # empty causal span; a stale seq_len of 1 overflows to non-finite KV. + self.input_buffers.seq_lens_buf[:bs].fill_(self.max_tokens_per_req) + self._init_capture_metadata(bs) + run_once() + + # Clear any per-pool state that warm-up dirtied at pool row 0, + # so the graph captures reads against a clean baseline. + if self.sampling_backend is not None: + self.sampling_backend.reset_capture_state() + + torch.cuda.synchronize() + dist.barrier() + + # Warmups can switch a backend back to eager metadata objects. Restore + # the graph-backed metadata immediately before capture so replay-time + # metadata refreshes update the same tensors recorded by the graph. + self._init_capture_metadata(bs) + + # Fill sampler buffers OUTSIDE the capture so RNG ops aren't recorded. + self._prepare_sampling_capture(bs=bs, variant=variant) + # Warmup forwards can mutate aliased metadata buffers, so refresh + # them again immediately before graph capture records the final views. + self._init_capture_metadata(bs) + + self.deepep_adapter.capture() + + global _is_capture_mode + _is_capture_mode = True + global global_graph_memory_pool + with torch.cuda.graph(graph, pool=global_graph_memory_pool, stream=self.stream): + out = run_once() + + torch.cuda.synchronize() + dist.barrier() + _is_capture_mode = False + + # Graph capture records the hostfunc launches without invoking + # them, so the dummy run_once pushed stays queued — drain it, and + # reset prev_batch/current_batch so the first real replay's build + # doesn't advance the matcher from a stale warmup entry. + if self.capturable_grammar is not None: + while True: + try: + self.capturable_grammar.queue.get_nowait() + except queue.Empty: + break + self.capturable_grammar.reset_state() + + global_graph_memory_pool = graph.pool() + return graph, out + + def _capture_paged_cache_block_tables(self, bs: int, pool) -> dict | None: + specs = tuple(pool.paged_cache_group_specs) + if not specs: + return None + out = {} + for spec in specs: + max_pages = compute_max_logical_pages_for_capture( + spec, + max_context_len=( + self.max_tokens_per_req * self.max_bs + if self.context_len <= 0 + else self.context_len + ), + max_tokens_per_req=self.max_tokens_per_req, + overlap_schedule_depth=self.overlap_schedule_depth, + ) + out[str(spec.group_id)] = torch.zeros( + (bs, max_pages), + dtype=torch.int32, + device=self.device, + ) + return out + + def _flat_cache_group_ids(self, pool) -> tuple[str, ...]: + """Group ids for flat per-group CUDA-graph capture: real tables only + arrive at replay, so capture needs just the ids to allocate its + persistent per-group buffers.""" + if not getattr(self.attn_backend, "uses_flat_cache_groups", False): + return () + return tuple(str(spec.group_id) for spec in pool.paged_cache_group_specs) + + def _draft_flat_group_ids(self) -> tuple[str, ...]: + """The draft head shares the target full-attention group's page ids + (EAGLE writes its own pool tensors at the same indices), so the + drafter consumes exactly that group's table on the flat path.""" + if self.draft_attn_backend is None or not getattr( + self.draft_attn_backend, "uses_flat_cache_groups", False + ): + return () + return tuple( + str(spec.group_id) + for spec in self.token_to_kv_pool.paged_cache_group_specs + if spec.family != "state" and spec.retention == "full_history" + ) + + def _draft_flat_tables(self, flat_block_tables): + """Subset of the target's per-group tables the drafter consumes.""" + gids = self._draft_flat_group_ids() + if not gids or not flat_block_tables: + return None + subset = { + gid: flat_block_tables[gid] for gid in gids if gid in flat_block_tables + } + return subset or None + + def _init_capture_metadata(self, bs: int): + capture_kwargs = {} + if self.input_buffers.has_mamba: + capture_kwargs["mamba_pool_indices"] = ( + self.input_buffers.mamba_pool_indices_buf[:bs] + ) + if self.attn_backend.uses_paged_cache_groups: + paged_cache_block_tables = self._capture_paged_cache_block_tables( + bs, + self.token_to_kv_pool, + ) + if paged_cache_block_tables is not None: + capture_kwargs["paged_cache_block_tables"] = paged_cache_block_tables + if self.drafter is not None: + capture_kwargs["num_tokens"] = bs * self.max_tokens_per_req + flat_cache_group_ids = self._flat_cache_group_ids(self.token_to_kv_pool) + if flat_cache_group_ids: + capture_kwargs["flat_cache_group_ids"] = flat_cache_group_ids + self.attn_backend.init_forward_metadata_capture_cuda_graph( + bs, + self.input_buffers.req_pool_indices_buf[:bs], + self.input_buffers.seq_lens_buf[:bs], + ForwardMode.DECODE, + **capture_kwargs, + ) + if self.draft_attn_backend is not None: + draft_kwargs = {} + if ( + self.draft_token_to_kv_pool is not None + and self.draft_attn_backend.uses_paged_cache_groups + ): + draft_paged_cache_block_tables = self._capture_paged_cache_block_tables( + bs, + self.draft_token_to_kv_pool, + ) + if draft_paged_cache_block_tables is not None: + draft_kwargs["paged_cache_block_tables"] = ( + draft_paged_cache_block_tables + ) + draft_kwargs["num_tokens"] = bs * self.max_tokens_per_req + draft_flat_ids = self._draft_flat_group_ids() + if draft_flat_ids: + draft_kwargs["flat_cache_group_ids"] = draft_flat_ids + # Drafter mutates seq_lens_buf in place per step; backends alias. + self.draft_attn_backend.init_forward_metadata_capture_cuda_graph( + bs, + self.input_buffers.req_pool_indices_buf[:bs], + self.input_buffers.seq_lens_buf[:bs], + ForwardMode.DECODE, + **draft_kwargs, + ) + + def _idle_flat_block_tables(self, padded_bs: int) -> dict | None: + """Minimal per-group tables for the bs==0 idle replay: all rows are + dummy rows, so one column of page-0 entries per group is valid. + None when the pool publishes no groups.""" + specs = tuple(self.token_to_kv_pool.paged_cache_group_specs) + if not specs: + return None + table = torch.zeros((padded_bs, 1), dtype=torch.int32, device=self.device) + return {str(spec.group_id): table for spec in specs} + + @staticmethod + def _pad_block_tables_to_padded_bs( + block_tables: dict, + *, + actual_bs: int, + padded_bs: int, + pad_value: int = -1, + ) -> dict: + """Pad each table with dummy ROWS up to padded_bs. Flat passes + pad_value=0, radix/V4 keeps -1 — see the padding contract at the MHA + backend's replay guard (backends/mha.py). + """ + if padded_bs <= actual_bs: + return block_tables + out = {} + for key, table in block_tables.items(): + if not isinstance(table, torch.Tensor): + out[key] = table + continue + rows = int(table.shape[0]) + if rows == padded_bs: + out[key] = table + continue + out[key] = torch.nn.functional.pad( + table, + (0, 0, 0, padded_bs - rows), + value=pad_value, + ) + return out + + @staticmethod + def _pad_offsets_to_padded_bs( + base_offsets: dict, + *, + actual_bs: int, + padded_bs: int, + ) -> dict: + if padded_bs <= actual_bs: + return base_offsets + out = {} + for key, off in base_offsets.items(): + if not isinstance(off, torch.Tensor): + out[key] = off + continue + rows = int(off.shape[0]) + if rows == padded_bs: + out[key] = off + continue + # Base 0: padded rows have no real request; the paired padded + # table row is invalid (-1). + out[key] = torch.nn.functional.pad( + off, + (0, padded_bs - rows), + value=0, + ) + return out + + def _init_replay_metadata( + self, + padded_bs: int, + actual_bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + req_to_page: torch.Tensor, + forward_mode: ForwardMode, + **kwargs, + ): + """Graph-replay path — update persistent cuda-graph buffers in place.""" + paged_cache_block_tables = kwargs.pop("paged_cache_block_tables", None) + paged_cache_block_table_base_offsets = kwargs.pop( + "paged_cache_block_table_base_offsets", None + ) + flat_block_tables = kwargs.pop("flat_block_tables", None) + target_uses_paged_groups = getattr( + self.attn_backend, + "uses_paged_cache_groups", + False, + ) + draft_uses_paged_groups = self.draft_attn_backend is not None and getattr( + self.draft_attn_backend, "uses_paged_cache_groups", False + ) + if paged_cache_block_tables is not None and ( + target_uses_paged_groups or draft_uses_paged_groups + ): + table_bs = next( + ( + int(table.shape[0]) + for table in paged_cache_block_tables.values() + if isinstance(table, torch.Tensor) + ), + int(req_pool_indices.shape[0]), + ) + paged_cache_block_tables = self._pad_block_tables_to_padded_bs( + paged_cache_block_tables, + actual_bs=table_bs, + padded_bs=padded_bs, + ) + if paged_cache_block_table_base_offsets is not None: + paged_cache_block_table_base_offsets = self._pad_offsets_to_padded_bs( + paged_cache_block_table_base_offsets, + actual_bs=actual_bs, + padded_bs=padded_bs, + ) + if target_uses_paged_groups: + kwargs["paged_cache_block_tables"] = paged_cache_block_tables + if paged_cache_block_table_base_offsets is not None: + kwargs["paged_cache_block_table_base_offsets"] = ( + paged_cache_block_table_base_offsets + ) + if flat_block_tables is not None and getattr( + self.attn_backend, "uses_flat_cache_groups", False + ): + flat_table_bs = next( + ( + int(table.shape[0]) + for table in flat_block_tables.values() + if isinstance(table, torch.Tensor) + ), + int(req_pool_indices.shape[0]), + ) + kwargs["flat_block_tables"] = self._pad_block_tables_to_padded_bs( + flat_block_tables, + actual_bs=flat_table_bs, + padded_bs=padded_bs, + pad_value=0, + ) + if self.attn_backend.uses_padded_decode_token_mask: + kwargs["actual_bs"] = actual_bs + if target_uses_paged_groups and getattr(self, "drafter", None) is not None: + kwargs["num_tokens"] = padded_bs * self.max_tokens_per_req + self.attn_backend.init_forward_metadata_replay_cuda_graph( + padded_bs, + req_pool_indices, + seq_lens, + req_to_page=req_to_page, + forward_mode=forward_mode, + **kwargs, + ) + if self.draft_attn_backend is not None: + draft_attn_kwargs = {} + if draft_uses_paged_groups and paged_cache_block_tables is not None: + draft_attn_kwargs["paged_cache_block_tables"] = paged_cache_block_tables + if paged_cache_block_table_base_offsets is not None: + draft_attn_kwargs["paged_cache_block_table_base_offsets"] = ( + paged_cache_block_table_base_offsets + ) + if getattr(self.draft_attn_backend, "uses_padded_decode_token_mask", False): + draft_attn_kwargs["actual_bs"] = actual_bs + draft_flat = self._draft_flat_tables(kwargs.get("flat_block_tables")) + if draft_flat is not None: + draft_attn_kwargs["flat_block_tables"] = draft_flat + draft_forward_mode = ForwardMode.DECODE + if draft_uses_paged_groups: + draft_attn_kwargs["num_tokens"] = padded_bs * self.max_tokens_per_req + draft_seq_lens = self.drafter.draft_seq_lens_buf[:padded_bs] + draft_seq_lens.copy_(seq_lens[:padded_bs]) + self.draft_attn_backend.init_forward_metadata_replay_cuda_graph( + padded_bs, + req_pool_indices, + draft_seq_lens, + req_to_page=self.drafter.req_to_page, + forward_mode=draft_forward_mode, + **draft_attn_kwargs, + ) + + @nvtx_range("attn_meta_prep", color="orange") + def _init_forward_metadata( + self, + padded_bs: int, + num_extends: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + req_to_page: torch.Tensor, + forward_mode: ForwardMode, + **kwargs, + ): + """Eager path — allocate/refresh metadata for the upcoming forward.""" + if ( + getattr(self.attn_backend, "uses_paged_cache_groups", False) + and self.drafter is not None + and forward_mode.is_decode() + ): + kwargs.setdefault("num_tokens", padded_bs * self.max_tokens_per_req) + self.attn_backend.init_forward_metadata( + bs=padded_bs, + num_extends=num_extends, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + req_to_page=req_to_page, + forward_mode=forward_mode, + **kwargs, + ) + if self.draft_attn_backend is not None: + draft_kwargs = {} + if getattr(self.draft_attn_backend, "uses_paged_cache_groups", False): + for key in ( + "paged_cache_block_tables", + "paged_cache_block_table_base_offsets", + ): + value = kwargs.get(key) + if value is not None: + draft_kwargs[key] = value + draft_flat = self._draft_flat_tables(kwargs.get("flat_block_tables")) + if draft_flat is not None: + draft_kwargs["flat_block_tables"] = draft_flat + + # The drafter mutates draft_seq_lens_buf between MTP draft steps; + # decode metadata must alias that buffer. + draft_seq_lens = self.drafter.draft_seq_lens_buf[:padded_bs] + draft_seq_lens.copy_(seq_lens[:padded_bs]) + if forward_mode.is_extend_or_mixed(): + # Non-V4 draft backends follow the legacy contract: a single + # EXTEND/MIXED metadata init fills both first-step prefill + # metadata and step 1+ decode metadata, with seq_lens aliased + # to the drafter-owned mutable buffer. V4 additionally needs + # the accepted-prefix view for first-step grouped-cache + # metadata, then a separate decode init to prepare the draft + # decode metadata from that first-step state. + draft_prefill_seq_lens = ( + seq_lens if self.use_v4_mtp_paged_metadata else draft_seq_lens + ) + # Drafter consumes only the full group's table (see _draft_flat_tables). + draft_extend_kwargs = ( + {**kwargs, "flat_block_tables": draft_flat} + if kwargs.get("flat_block_tables") is not None + else kwargs + ) + self.draft_attn_backend.init_forward_metadata( + bs=padded_bs, + num_extends=num_extends, + req_pool_indices=req_pool_indices, + seq_lens=draft_prefill_seq_lens, + req_to_page=self.drafter.req_to_page, + forward_mode=forward_mode, + **draft_extend_kwargs, + ) + if self.use_v4_mtp_paged_metadata: + self.draft_attn_backend.init_forward_metadata( + bs=padded_bs, + num_extends=0, + req_pool_indices=req_pool_indices, + seq_lens=draft_seq_lens, + req_to_page=self.drafter.req_to_page, + forward_mode=ForwardMode.DECODE, + **draft_kwargs, + ) + else: + draft_metadata_seq_lens = ( + seq_lens if self.use_v4_mtp_paged_metadata else draft_seq_lens + ) + draft_forward_mode = ForwardMode.DECODE + if getattr(self.draft_attn_backend, "uses_paged_cache_groups", False): + draft_kwargs["num_tokens"] = padded_bs * self.max_tokens_per_req + self.draft_attn_backend.init_forward_metadata( + bs=padded_bs, + num_extends=0, + req_pool_indices=req_pool_indices, + seq_lens=draft_metadata_seq_lens, + req_to_page=self.drafter.req_to_page, + forward_mode=draft_forward_mode, + **draft_kwargs, + ) + + def _global_graph_bs(self, ctx: ForwardContext) -> int | None: + if self.dp_size <= 1 or ctx.global_num_tokens is None: + return None + max_num_tokens = max(ctx.global_num_tokens) + return (max_num_tokens + self.max_tokens_per_req - 1) // self.max_tokens_per_req + + def _can_use_graph(self, bs: int, ctx: ForwardContext) -> bool: + if self.disable: + return False + if not ctx.forward_mode.is_decode(): + return False + if self.dp_size > 1: + if not ctx.all_decode_or_idle: + return False + global_bs = self._global_graph_bs(ctx) + if global_bs is None or global_bs == 0: + return False + if self.disable_padding: + return self._has_cuda_graph_for_bs(global_bs) + return global_bs <= self.max_bs + if self.disable_padding: + return self._has_cuda_graph_for_bs(bs) + return bs <= self.max_bs + + def can_run(self, bs: int, ctx: ForwardContext) -> bool: + return self._can_use_graph(bs, ctx) + + def padded_bs(self, bs: int, ctx: ForwardContext) -> int: + return self._padded_bs(bs, ctx) + + def _padded_bs(self, bs: int, ctx: ForwardContext) -> int: + graph_bs = self._global_graph_bs(ctx) + target_bs = graph_bs if graph_bs is not None else bs + index = bisect.bisect_left(self.capture_bs, target_bs) + return self.capture_bs[index] + + def _pad_graph_req_pool_indices( + self, active_req_pool_indices: torch.Tensor, padded_bs: int + ) -> torch.Tensor: + pad = padded_bs - active_req_pool_indices.shape[0] + if pad <= 0: + return active_req_pool_indices + if self.config.spec_algo == "DFLASH": + # Route padding rows to the sentinel req-pool slot + # (max_req_pool_size), not slot 0. The DFLASH draft derives each + # row's block seq_len from valid_cache_lengths[req_pool], so + # padding rows pointing at slot 0 would grow unbounded with + # request 0's context and hang the draft block-decode kernel. + # The sentinel row stays zero-init (length 0, dummy page 0). + sentinel = int(self.config.max_req_pool_size) + return torch.cat( + [ + active_req_pool_indices, + active_req_pool_indices.new_full((pad,), sentinel), + ] + ) + return torch.cat( + [active_req_pool_indices, active_req_pool_indices.new_zeros(pad)] + ) + + def _set_graph_state_write_indices( + self, active_req_pool_indices: torch.Tensor, padded_bs: int + ) -> None: + state_indices = self.input_buffers.state_write_req_pool_indices_buf[:padded_bs] + active_bs = active_req_pool_indices.shape[0] + if active_bs > 0: + state_indices[:active_bs].copy_(active_req_pool_indices) + if active_bs < padded_bs: + state_indices[active_bs:padded_bs].fill_(int(self.config.max_req_pool_size)) + + def __call__( + self, + bs: int, + ctx: ForwardContext, + sampling_info: SamplingBatchInfo, + req_to_page: torch.Tensor, + extend_with_prefix: bool = False, + extend_prefix_lens: torch.Tensor | None = None, + extend_prefix_lens_cpu: torch.Tensor | None = None, + extend_seq_lens: torch.Tensor | None = None, + extend_seq_lens_cpu: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + out_cache_loc: torch.Tensor | None = None, + mamba_pool_indices: torch.Tensor | None = None, + mamba_cow_src_indices: torch.Tensor | None = None, + mamba_branching_seqlens: torch.Tensor | None = None, + mamba_track_pool_indices: torch.Tensor | None = None, + spec_info=None, + paged_cache_block_tables: dict | None = None, + paged_cache_block_table_base_offsets: dict | None = None, + flat_block_tables: dict | None = None, + ): + """ + Unified forward entry point. + + Dispatches to the captured CUDA graph when possible; falls back to the + eager forward_func otherwise. The caller does not need to know which + path was taken. + """ + use_graph = self._can_use_graph(bs, ctx) + padded_bs = self._padded_bs(bs, ctx) if use_graph else bs + active_req_pool_indices = self.input_buffers.req_pool_indices_buf[:bs] + + if use_graph and padded_bs != bs: + ctx.bs = padded_bs + pad = padded_bs - bs + seq_lens = torch.nn.functional.pad( + self.input_buffers.seq_lens_buf[:bs], (0, pad), value=1 + ) + req_pool_indices = self._pad_graph_req_pool_indices( + active_req_pool_indices, padded_bs + ) + self.input_buffers.seq_lens_buf[:padded_bs].copy_(seq_lens) + self.input_buffers.req_pool_indices_buf[:padded_bs].copy_(req_pool_indices) + if mamba_pool_indices is not None: + # Pad with -1 (PAD_SLOT_ID), NOT 0. Mamba slot 0 is a real + # allocatable slot, so padding with 0 aliases a live request's + # mamba state and corrupts it. -1 is the kernel-skipped pad id. + mamba_pool_indices = torch.nn.functional.pad( + mamba_pool_indices, (0, pad), value=-1 + ) + if mamba_cow_src_indices is not None: + mamba_cow_src_indices = torch.nn.functional.pad( + mamba_cow_src_indices, (0, pad), value=-1 + ) + if mamba_branching_seqlens is not None: + mamba_branching_seqlens = torch.nn.functional.pad( + mamba_branching_seqlens, (0, pad), value=-1 + ) + if mamba_track_pool_indices is not None: + mamba_track_pool_indices = torch.nn.functional.pad( + mamba_track_pool_indices, (0, pad), value=-1 + ) + else: + seq_lens = self.input_buffers.seq_lens_buf[:padded_bs] + req_pool_indices = self.input_buffers.req_pool_indices_buf[:padded_bs] + + if use_graph: + self._set_graph_state_write_indices(active_req_pool_indices, padded_bs) + + mamba_kwargs = {} + if mamba_pool_indices is not None: + mamba_kwargs["mamba_pool_indices"] = mamba_pool_indices + if mamba_cow_src_indices is not None: + mamba_kwargs["mamba_cow_src_indices"] = mamba_cow_src_indices + if mamba_branching_seqlens is not None: + mamba_kwargs["mamba_branching_seqlens"] = mamba_branching_seqlens + if mamba_track_pool_indices is not None: + mamba_kwargs["mamba_track_pool_indices"] = mamba_track_pool_indices + + if use_graph: + if ( + bs == 0 + and paged_cache_block_tables is None + and self.attn_backend.uses_paged_cache_groups + ): + paged_cache_block_tables = self._capture_paged_cache_block_tables( + padded_bs, + self.token_to_kv_pool, + ) + # The backend's stale-table guard also covers the bs==0 idle + # replay: synthesize minimal valid tables for it. + if ( + bs == 0 + and not flat_block_tables + and getattr(self.attn_backend, "uses_flat_cache_groups", False) + ): + flat_block_tables = self._idle_flat_block_tables(padded_bs) + self._init_replay_metadata( + padded_bs, + bs, + req_pool_indices, + seq_lens, + req_to_page=req_to_page, + forward_mode=ctx.forward_mode, + num_padding=padded_bs - bs if padded_bs != bs else 0, + paged_cache_block_tables=paged_cache_block_tables, + paged_cache_block_table_base_offsets=( + paged_cache_block_table_base_offsets + ), + flat_block_tables=flat_block_tables, + **mamba_kwargs, + ) + + # Runtime prepare() is called by ModelExecutor with per-request rids + # BEFORE self.forward_step — we don't refill here to avoid clobbering + # the per-request generators with the capture-stub generator. + self.deepep_adapter.replay() + + graph_key = self._cuda_graph_key(padded_bs) + with nvtx_range("graph_replay", color="red"): + self.graphs[graph_key].replay() + + ( + output_tokens, + output_lengths, + output_logprobs, + ) = self.output_buffers[graph_key] + + result = ( + output_tokens[: bs * self.max_tokens_per_req], + output_lengths[:bs], + ( + output_logprobs[: bs * self.max_tokens_per_req] + if output_logprobs is not None + else None + ), + ) + + else: + # Eager parity with the replay stale-table guard: with >1 group + # the single-table fallback would serve first-group pages to + # every layer. Idle/bs==0 forwards carry no requests (exempt); + # a single published group falls back to the single table. + if ( + bs > 0 + and not ctx.forward_mode.is_idle() + and not flat_block_tables + and getattr(self.attn_backend, "uses_flat_cache_groups", False) + and len(self.token_to_kv_pool.paged_cache_group_specs) > 1 + ): + raise RuntimeError( + "CudaGraphWrapper eager forward: pool publishes " + f"{len(self.token_to_kv_pool.paged_cache_group_specs)} " + "flat cache groups and the backend consumes flat tables, " + f"but flat_block_tables is missing/empty at bs={bs} " + f"({ctx.forward_mode.name}); the single-table fallback " + "would use one group's pages for all layers." + ) + metadata_num_tokens = ( + {"num_tokens": ctx.input_num_tokens} + if self.attn_backend.uses_paged_cache_groups + else {} + ) + self._init_forward_metadata( + padded_bs, + ctx.num_extends, + req_pool_indices, + seq_lens, + req_to_page=req_to_page, + forward_mode=ctx.forward_mode, + extend_with_prefix=extend_with_prefix, + extend_prefix_lens=extend_prefix_lens, + extend_prefix_lens_cpu=extend_prefix_lens_cpu, + extend_seq_lens=extend_seq_lens, + extend_seq_lens_cpu=extend_seq_lens_cpu, + positions=positions, + out_cache_loc=out_cache_loc, + global_num_tokens=ctx.global_num_tokens, + all_decode_or_idle=ctx.all_decode_or_idle, + capture_hidden_mode=ctx.capture_hidden_mode, + spec_info=spec_info, + **metadata_num_tokens, + paged_cache_block_tables=( + paged_cache_block_tables + if self.attn_backend.uses_paged_cache_groups + else None + ), + paged_cache_block_table_base_offsets=( + paged_cache_block_table_base_offsets + if self.attn_backend.uses_paged_cache_groups + else None + ), + flat_block_tables=( + flat_block_tables + if self.attn_backend.uses_flat_cache_groups + else None + ), + **mamba_kwargs, + ) + + result = self._forward_func(bs=bs, ctx=ctx, sampling_info=sampling_info) + + if use_graph and padded_bs != bs: + ctx.bs = bs + + # Update mamba/GDN state after speculative verify + if _should_update_mamba_state_after_mtp_verify( + self.drafter, self.attn_backend, ctx.forward_mode + ): + accept_lengths = result[1] + self.attn_backend.update_mamba_state_after_mtp_verify(accept_lengths, None) + + return result diff --git a/python/tokenspeed/runtime/execution/distributed_initializer.py b/python/tokenspeed/runtime/execution/distributed_initializer.py new file mode 100644 index 0000000..9d2771a --- /dev/null +++ b/python/tokenspeed/runtime/execution/distributed_initializer.py @@ -0,0 +1,193 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from dataclasses import dataclass + +import torch + +from tokenspeed.runtime.distributed.process_group_manager import ( + process_group_manager as pg_manager, +) +from tokenspeed.runtime.utils import ( + get_available_gpu_memory, + get_colorful_logger, +) +from tokenspeed.runtime.utils.common import maybe_set_numa_aware_cpu_affinity +from tokenspeed.runtime.utils.server_args import PortArgs, ServerArgs + +logger = get_colorful_logger(__name__) + + +@dataclass +class DistributedConfig: + """Lightweight configuration for distributed initialization. + + Contains only primitive types (int, str, bool) to avoid heavy dependencies. + All information needed for distributed setup is captured here. + """ + + # Device configuration + device: str + gpu_id: int + + # Distributed topology + world_size: int + global_rank: int + local_rank: int + + # Tensor parallelism + attn_tp_rank: int + attn_tp_size: int + + # Data parallelism + dp_size: int + + # Dense layer parallelism + dense_tp_size: int + + # Expert parallelism (MoE) + moe_ep_size: int + moe_ep_rank: int + + # Network configuration + nccl_port: int + dist_init_addr: str | None = None + distributed_timeout_seconds: int = 1800 + + # Node configuration + nnodes: int = 1 + nprocs_per_node: int = 1 + + # Model configuration (needed for attention groups) + hidden_size: int = 0 + max_num_tokens: int = 0 + + # Feature flags + disable_custom_all_reduce: bool = False + force_deterministic_rsag: bool = False + + # The full Mapping object for pg_manager initialization + mapping: object = None + + @classmethod + def from_server_args( + cls, + server_args: ServerArgs, + port_args: PortArgs, + gpu_id: int, + global_rank: int, + hidden_size: int, + max_num_tokens: int, + ): + mapping = server_args.mapping + return cls( + device=server_args.device, + gpu_id=gpu_id, + world_size=mapping.world_size, + global_rank=global_rank, + local_rank=global_rank % mapping.nprocs_per_node, + attn_tp_rank=mapping.attn.tp_rank, + attn_tp_size=mapping.attn.tp_size, + dp_size=mapping.attn.dp_size, + dense_tp_size=mapping.dense.tp_size, + moe_ep_size=mapping.moe.ep_size, + moe_ep_rank=mapping.moe.ep_rank, + nccl_port=port_args.nccl_port, + dist_init_addr=server_args.dist_init_addr, + distributed_timeout_seconds=( + server_args.distributed_timeout_seconds + if server_args.distributed_timeout_seconds is not None + else 1800 + ), + nnodes=mapping.nnodes, + nprocs_per_node=mapping.nprocs_per_node, + hidden_size=hidden_size, + max_num_tokens=max_num_tokens, + disable_custom_all_reduce=server_args.disable_custom_all_reduce, + force_deterministic_rsag=server_args.force_deterministic_rsag, + mapping=mapping, + ) + + +class DistributedInitializer: + @staticmethod + def initialize(config: DistributedConfig) -> float: + torch.get_device_module(config.device).set_device(config.gpu_id) + logger.info( + "Init torch distributed begin. Avail mem=%.4f GB", + get_available_gpu_memory(config.device, config.gpu_id), + ) + if config.device == "cuda": + maybe_set_numa_aware_cpu_affinity(config.gpu_id) + + # Determine backend + if config.device == "cuda": + backend = "nccl" + else: + raise ValueError(f"Unsupported device: {config.device}") + + # Build distributed init method + if config.dist_init_addr: + dist_init_method = f"tcp://{config.dist_init_addr}" + else: + dist_init_method = f"tcp://127.0.0.1:{config.nccl_port}" + + # Initialize distributed via the mapping-based process group manager + pg_manager.init_distributed( + config.mapping, + backend=backend, + distributed_init_method=dist_init_method, + timeout=config.distributed_timeout_seconds, + ) + pg_manager.init_process_group(config.mapping.world_group) + pg_manager.init_process_group(config.mapping.attn.tp_group) + pg_manager.init_process_group(config.mapping.dense.tp_group) + pg_manager.init_process_group(config.mapping.moe.tp_ep_group) + + logger.info( + "Init comm buff end. Avail mem=%.4f GB", + get_available_gpu_memory(config.device, config.gpu_id), + ) + mapping = config.mapping + logger.info( + "Current Process distributed state: global rank: %s attn_tp_rank: %s attn_dp_rank: %s", + mapping.rank, + mapping.attn.tp_rank, + mapping.attn.dp_rank, + ) + + # Get minimum available GPU memory across all ranks + min_per_gpu_memory = get_available_gpu_memory( + config.device, + config.gpu_id, + distributed=config.world_size > 1, + cpu_group=pg_manager.get_process_group("gloo", mapping.world_group), + ) + + # Verify memory balance for tensor parallelism + if config.world_size > 1: + local_gpu_memory = get_available_gpu_memory(config.device, config.gpu_id) + if min_per_gpu_memory < local_gpu_memory * 0.9: + raise ValueError( + "The memory capacity is unbalanced. " + "Some GPUs may be occupied by other processes." + ) + + return min_per_gpu_memory diff --git a/python/tokenspeed/runtime/execution/drafter/__init__.py b/python/tokenspeed/runtime/execution/drafter/__init__.py new file mode 100644 index 0000000..2b589d8 --- /dev/null +++ b/python/tokenspeed/runtime/execution/drafter/__init__.py @@ -0,0 +1,19 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. diff --git a/python/tokenspeed/runtime/execution/drafter/base.py b/python/tokenspeed/runtime/execution/drafter/base.py new file mode 100644 index 0000000..13d88ba --- /dev/null +++ b/python/tokenspeed/runtime/execution/drafter/base.py @@ -0,0 +1,83 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from abc import abstractmethod +from typing import TYPE_CHECKING + +import torch + +from tokenspeed.runtime.execution.model_runner import ModelRunner + +if TYPE_CHECKING: + from tokenspeed.runtime.execution.context import ForwardContext + from tokenspeed.runtime.execution.input_buffer import InputBuffers + from tokenspeed.runtime.execution.runtime_states import RuntimeStates + from tokenspeed.runtime.layers.attention.backends.base import AttentionBackend + from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool + from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput + + +class BaseDrafter: + def __init__( + self, + spec_num_tokens: int, + spec_num_steps: int | None = None, + draft_model_runner: ModelRunner | None = None, + runtime_states: RuntimeStates | None = None, + input_buffers: InputBuffers | None = None, + page_size: int | None = None, + req_to_page: torch.Tensor | None = None, + attn_backend: AttentionBackend | None = None, + token_to_kv_pool: BaseTokenToKVPool | None = None, + vocab_size: int | None = None, + ): + self.spec_num_tokens = spec_num_tokens + self.spec_num_steps = spec_num_steps + self.draft_model_runner = draft_model_runner + self.runtime_states = runtime_states + self.input_buffers = input_buffers + self.page_size = page_size + self.req_to_page = req_to_page + self.attn_backend = attn_backend + self.token_to_kv_pool = token_to_kv_pool + self.vocab_size = vocab_size + + @abstractmethod + def get_candidates( + self, + base_ctx: ForwardContext, + ) -> torch.Tensor | None: + raise NotImplementedError + + @abstractmethod + def run( + self, + base_ctx: ForwardContext, + logits_output: LogitsProcessorOutput, + output_tokens: torch.Tensor, + accept_lengths: torch.Tensor, + ) -> torch.Tensor: + raise NotImplementedError + + @abstractmethod + def draft(self, *args, **kwargs) -> torch.Tensor | None: + raise NotImplementedError diff --git a/python/tokenspeed/runtime/execution/drafter/dflash.py b/python/tokenspeed/runtime/execution/drafter/dflash.py new file mode 100644 index 0000000..c62d52d --- /dev/null +++ b/python/tokenspeed/runtime/execution/drafter/dflash.py @@ -0,0 +1,569 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from tokenspeed.runtime.distributed.comm_ops import all_gather_into_tensor +from tokenspeed.runtime.execution.cache_loc_kernel import compute_out_cache_loc_uniform +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.execution.drafter.base import BaseDrafter +from tokenspeed.runtime.execution.forward_batch_info import ( + CaptureHiddenMode, + ForwardMode, +) +from tokenspeed.runtime.layers.logits_processor import LogitsMetadata +from tokenspeed.runtime.utils import get_colorful_logger +from tokenspeed.runtime.utils.nvtx import nvtx_range + +if TYPE_CHECKING: + from tokenspeed.runtime.execution.input_buffer import InputBuffers + from tokenspeed.runtime.execution.model_runner import ModelRunner + from tokenspeed.runtime.execution.runtime_states import RuntimeStates + from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput + +logger = get_colorful_logger(__name__) + + +class DFlash(BaseDrafter): + """DFlash block drafter backed by a native TokenSpeed draft model.""" + + def __init__( + self, + spec_num_tokens: int, + spec_num_steps: int, + page_size: int, + draft_model_runner: ModelRunner | None = None, + req_to_page: torch.Tensor | None = None, + attn_backend=None, + token_to_kv_pool=None, + runtime_states: RuntimeStates | None = None, + input_buffers: InputBuffers | None = None, + vocab_size: int | None = None, + ) -> None: + super().__init__( + spec_num_tokens=spec_num_tokens, + spec_num_steps=spec_num_steps, + draft_model_runner=draft_model_runner, + runtime_states=runtime_states, + input_buffers=input_buffers, + page_size=page_size, + req_to_page=req_to_page, + attn_backend=attn_backend, + token_to_kv_pool=token_to_kv_pool, + vocab_size=vocab_size, + ) + if draft_model_runner is None: + raise ValueError("Native DFLASH requires a draft model runner.") + + server_args = draft_model_runner.server_args + if not server_args.speculative_draft_model_path: + raise ValueError("DFLASH requires --speculative-draft-model-path.") + + self.device = torch.device(draft_model_runner.device) + self.model = draft_model_runner.model + + cfg = self.model.config + dflash_cfg = getattr(cfg, "dflash_config", {}) or {} + self.target_layer_ids = [int(x) for x in dflash_cfg.get("target_layer_ids", [])] + if not self.target_layer_ids: + raise ValueError( + "DFLASH draft config must define dflash_config.target_layer_ids." + ) + if "mask_token_id" not in dflash_cfg: + raise ValueError( + "DFLASH draft config must define dflash_config.mask_token_id." + ) + self.mask_token_id = int(dflash_cfg["mask_token_id"]) + self.block_size = int(getattr(cfg, "block_size", spec_num_tokens)) + if self.block_size != int(spec_num_tokens): + logger.warning( + "DFLASH block size mismatch: checkpoint block_size=%s, " + "runtime speculative_num_draft_tokens=%s.", + self.block_size, + spec_num_tokens, + ) + self.hidden_size = int(getattr(cfg, "hidden_size")) + self.idle_forward_steps = 1 + self._init_native_buffers() + self._greedy_gathered_max: torch.Tensor | None = None + self._greedy_gathered_ids: torch.Tensor | None = None + self._greedy_gather_cap = 0 + + def _init_native_buffers(self) -> None: + if self.input_buffers is None: + raise ValueError("Native DFLASH requires input buffers.") + if self.req_to_page is None: + raise ValueError("Native DFLASH requires req_to_page.") + if self.attn_backend is None or self.token_to_kv_pool is None: + raise ValueError("Native DFLASH requires draft attention components.") + + max_bs = self.input_buffers.max_bs + self.draft_seq_lens_buf = torch.zeros_like(self.input_buffers.seq_lens_buf) + self.draft_out_cache_loc_buf = torch.empty( + (max_bs * self.spec_num_tokens,), + dtype=torch.int32, + device=self.device, + ) + self.draft_input_lengths_buf = torch.full( + (max_bs,), + self.spec_num_tokens, + dtype=torch.int32, + device=self.device, + ) + self.draft_extend_seq_lens_cpu = torch.full( + (max_bs,), + self.spec_num_tokens, + dtype=torch.int32, + pin_memory=True, + ) + self.block_offsets = torch.arange( + self.spec_num_tokens, dtype=torch.int64, device=self.device + ) + self.block_ids_buf = torch.empty( + (max_bs, self.spec_num_tokens), dtype=torch.int32, device=self.device + ) + self.block_positions_buf = torch.empty( + (max_bs, self.spec_num_tokens), dtype=torch.int64, device=self.device + ) + + def bind_target_model(self, target_model) -> None: + language_model = getattr(target_model, "language_model", target_model) + self.target_model = target_model + self.target_language_model = language_model + self.embed_tokens = target_model.get_input_embeddings() + self.lm_head = target_model.lm_head + self.logits_processor = language_model.logits_processor + + def _greedy_gather_capacity(self) -> int: + """Max element count for the greedy head's tensor-parallel all-gather + scratch: a full ``max_bs`` decode block. + + The greedy head samples the last ``spec_num_tokens - 1`` block + positions per request and all-gathers them across the TP group, so the + worst case is ``tp_size * max_bs * (spec_num_tokens - 1)``. + """ + tp_size = int(self.logits_processor.tp_size) + return tp_size * self.input_buffers.max_bs * max(self.spec_num_tokens - 1, 1) + + def _ensure_greedy_gather_buffers( + self, + max_dtype: torch.dtype, + ids_dtype: torch.dtype, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Lazily create the greedy all-gather scratch ONCE at its maximum + capacity, then reuse it in place for every batch size. + + Sizing to the max ``max_bs`` block (rather than growing per batch size) + is required for CUDA-graph correctness. Graphs are captured for + increasing batch sizes (``[1, 2, ..., max_bs]``); a buffer grown lazily + would be freed and reallocated when a larger bs needs more room, leaving + every smaller-bs graph captured earlier with an + ``all_gather_into_tensor`` recorded against freed memory. On replay + those small-bs decode steps read garbage (out-of-vocab) draft token ids, + which flow into the next verify forward's embedding lookup and trigger a + CUDA illegal memory access. A fixed max-capacity buffer is allocated + during warmup (before capture) and shared by every captured graph. + + Returns the (max, id) scratch tensors; callers slice ``[:needed]``. + """ + cap = self._greedy_gather_capacity() + if ( + self._greedy_gathered_max is None + or self._greedy_gathered_ids is None + or self._greedy_gather_cap < cap + or self._greedy_gathered_max.dtype != max_dtype + or self._greedy_gathered_max.device != device + or self._greedy_gathered_ids.dtype != ids_dtype + ): + self._greedy_gathered_max = torch.empty( + (cap,), dtype=max_dtype, device=device + ) + self._greedy_gathered_ids = torch.empty( + (cap,), dtype=ids_dtype, device=device + ) + self._greedy_gather_cap = cap + return self._greedy_gathered_max, self._greedy_gathered_ids + + def _greedy_sample_from_vocab_parallel_head( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + if not hasattr(self.lm_head, "weight") or not hasattr( + self.lm_head, "shard_indices" + ): + metadata = LogitsMetadata(forward_mode=ForwardMode.DECODE) + logits = self.logits_processor._get_logits( + hidden_states, self.lm_head, metadata + ) + return torch.argmax(logits, dim=-1).to(torch.int32) + + shard = self.lm_head.shard_indices + weight = self.lm_head.weight + hidden_states = hidden_states.to(weight.dtype) + + num_org = int(shard.num_org_elements) + num_org_padded = int(shard.num_org_elements_padded) + num_added = int(shard.num_added_elements) + org_vocab_start = int(shard.org_vocab_start_index) + added_vocab_start = int(shard.added_vocab_start_index) + + chunk_len = int(hidden_states.shape[0]) + if num_org > 0: + base_logits = torch.matmul(hidden_states, weight[:num_org].T) + local_max, local_arg = torch.max(base_logits, dim=-1) + else: + local_max = torch.full( + (chunk_len,), + torch.finfo(weight.dtype).min, + dtype=weight.dtype, + device=hidden_states.device, + ) + local_arg = torch.zeros( + (chunk_len,), dtype=torch.int64, device=hidden_states.device + ) + + if num_added > 0: + added_start = num_org_padded + added_end = num_org_padded + num_added + added_weight = weight[added_start:added_end] + added_logits = torch.matmul(hidden_states, added_weight.T) + added_max, added_arg = torch.max(added_logits, dim=-1) + use_added = added_max > local_max + local_max = torch.where(use_added, added_max, local_max) + local_arg = torch.where( + use_added, + added_arg.to(local_arg.dtype) + num_org_padded, + local_arg, + ) + + if num_added == 0: + global_ids = local_arg + org_vocab_start + else: + global_ids = torch.empty( + (chunk_len,), dtype=torch.int64, device=hidden_states.device + ) + is_base = local_arg < num_org + global_ids[is_base] = org_vocab_start + local_arg[is_base] + global_ids[~is_base] = added_vocab_start + ( + local_arg[~is_base] - num_org_padded + ) + + tp_size = int(self.logits_processor.tp_size) + if tp_size == 1: + return global_ids.to(torch.int32) + + needed = tp_size * chunk_len + gathered_max, gathered_ids = self._ensure_greedy_gather_buffers( + local_max.dtype, global_ids.dtype, hidden_states.device + ) + gathered_max = gathered_max[:needed] + gathered_ids = gathered_ids[:needed] + all_gather_into_tensor( + gathered_max, + local_max.contiguous(), + self.logits_processor.tp_group, + ) + all_gather_into_tensor( + gathered_ids, + global_ids.contiguous(), + self.logits_processor.tp_group, + ) + + gathered_max = gathered_max.view(tp_size, chunk_len) + gathered_ids = gathered_ids.view(tp_size, chunk_len) + best_rank = torch.argmax(gathered_max, dim=0).unsqueeze(0) + return torch.gather(gathered_ids, 0, best_rank).view(-1).to(torch.int32) + + @nvtx_range("dflash_update_native_cache", color="purple") + def _update_native_cache_from_target( + self, + base_ctx: ForwardContext, + logits_output: LogitsProcessorOutput, + accept_lengths: torch.Tensor, + ) -> None: + hidden = logits_output.hidden_states + if hidden is None: + raise RuntimeError("DFLASH requires target hidden states.") + if hidden.shape[0] != base_ctx.input_num_tokens: + raise RuntimeError( + "DFLASH hidden-state/token mismatch: " + f"hidden_tokens={hidden.shape[0]}, input_tokens={base_ctx.input_num_tokens}." + ) + + bs = base_ctx.bs + # The target verify forward emits spec_num_tokens hidden states per + # decode request (the candidate block); input_lengths_buf only tracks + # the committed-token count there, so split decode rows by + # spec_num_tokens. Prefill rows keep their real chunk lengths. + lengths = self.input_buffers.input_lengths_buf[:bs].to(torch.int64).clone() + lengths[base_ctx.num_extends :] = self.spec_num_tokens + req_pool_indices = self.input_buffers.req_pool_indices_buf[:bs] + positions = self.input_buffers.positions_buf[: base_ctx.input_num_tokens] + cache_locs = self.input_buffers.out_cache_loc_buf[: base_ctx.input_num_tokens] + + if ( + base_ctx.num_extends == 0 + and torch.cuda.is_available() + and torch.cuda.is_current_stream_capturing() + ): + old_lens = self.runtime_states.valid_cache_lengths.index_select( + 0, req_pool_indices + ) + self.draft_seq_lens_buf[:bs].copy_( + old_lens.to(torch.int32) + accept_lengths[:bs].to(torch.int32) + ) + self._write_native_cache(hidden, positions, cache_locs) + return + + hidden_chunks = torch.split(hidden, lengths.detach().cpu().tolist(), dim=0) + pos_chunks = torch.split(positions, lengths.detach().cpu().tolist(), dim=0) + loc_chunks = torch.split(cache_locs, lengths.detach().cpu().tolist(), dim=0) + + selected_hidden = [] + selected_positions = [] + selected_cache_locs = [] + new_seq_lens = torch.empty((bs,), dtype=torch.int32, device=self.device) + + for row, (chunk, pos_chunk, loc_chunk) in enumerate( + zip(hidden_chunks, pos_chunks, loc_chunks, strict=True) + ): + if row < base_ctx.num_extends: + take = int(chunk.shape[0]) + else: + take = int(accept_lengths[row].item()) + if take <= 0: + pool_idx = req_pool_indices[row] + new_seq_lens[row] = self.runtime_states.valid_cache_lengths[pool_idx] + continue + + chunk = chunk[:take].contiguous() + pos_chunk = pos_chunk[:take].contiguous() + loc_chunk = loc_chunk[:take].contiguous() + selected_hidden.append(chunk) + selected_positions.append(pos_chunk) + selected_cache_locs.append(loc_chunk) + new_seq_lens[row] = (pos_chunk[-1] + 1).to(torch.int32) + + self.draft_seq_lens_buf[:bs].copy_(new_seq_lens) + if not selected_hidden: + return + + target_hidden = torch.cat(selected_hidden, dim=0) + target_positions = torch.cat(selected_positions, dim=0) + target_cache_locs = torch.cat(selected_cache_locs, dim=0) + self._write_native_cache(target_hidden, target_positions, target_cache_locs) + + def _write_native_cache( + self, + target_hidden: torch.Tensor, + target_positions: torch.Tensor, + target_cache_locs: torch.Tensor, + ) -> None: + target_hidden = target_hidden.to( + device=self.device, + dtype=self.draft_model_runner.model.fc.weight.dtype, + ) + expected_width = int(self.draft_model_runner.model.fc.in_features) + actual_width = int(target_hidden.shape[-1]) + if actual_width != expected_width: + raise RuntimeError( + "DFLASH captured hidden width mismatch: " + f"expected {expected_width}, got {actual_width}. " + "Check dflash_config.target_layer_ids against the target model." + ) + with torch.inference_mode(): + ctx_hidden = self.draft_model_runner.model.project_target_hidden( + target_hidden + ) + for layer in self.draft_model_runner.model.layers: + attn = layer.self_attn + k, v = attn.kv_proj_only(ctx_hidden) + k = attn.apply_k_norm(k) + k = attn.apply_k_rope(target_positions, k) + k = k.view(-1, attn.num_kv_heads, attn.head_dim) + v = v.view(-1, attn.num_kv_heads, attn.head_dim) + self.token_to_kv_pool.set_kv_buffer( + attn.attn, + target_cache_locs, + k, + v, + attn.attn.k_scale, + attn.attn.v_scale, + ) + + @staticmethod + def _current_tokens_from_output( + output_tokens: torch.Tensor, + accept_lengths: torch.Tensor, + num_extends: int, + spec_num_tokens: int, + ) -> torch.Tensor: + bs = accept_lengths.shape[0] + current = torch.empty((bs,), dtype=torch.int32, device=output_tokens.device) + if num_extends > 0: + current[:num_extends] = output_tokens[:num_extends] + num_decodes = bs - num_extends + if num_decodes > 0: + offsets = ( + torch.arange( + num_decodes, dtype=torch.int64, device=output_tokens.device + ) + * spec_num_tokens + - 1 + + num_extends + ) + # ``accept_lengths`` can be clamped to 0 at the context limit. The + # request will be finished by the scheduler, but the drafter still + # runs for graph shape. Select a valid in-row dummy token instead of + # producing ``row * N - 1`` or crossing into the previous row. + safe_accept_lengths = ( + accept_lengths[num_extends:].to(torch.int64).clamp(1, spec_num_tokens) + ) + current[num_extends:] = output_tokens[offsets + safe_accept_lengths] + return current + + def get_candidates(self, base_ctx: ForwardContext) -> torch.Tensor | None: + num_extends = base_ctx.num_extends + num_decodes = base_ctx.bs - num_extends + if num_decodes == 0: + return None + num_decode_tokens = num_decodes * self.spec_num_tokens + num_prefill_tokens = base_ctx.input_num_tokens - num_decode_tokens + return self.input_buffers.input_ids_buf[ + num_prefill_tokens : base_ctx.input_num_tokens + ].reshape(num_decodes, self.spec_num_tokens) + + def draft(self, current_tokens: torch.Tensor) -> torch.Tensor: + return self._draft_native(current_tokens) + + @nvtx_range("dflash_native_draft", color="purple") + def _draft_native(self, current_tokens: torch.Tensor) -> torch.Tensor: + bs = current_tokens.shape[0] + req_pool_indices = self.input_buffers.req_pool_indices_buf[:bs] + prefix_lens = self.draft_seq_lens_buf[:bs].clone() + seq_lens_after = self.draft_seq_lens_buf[:bs] + seq_lens_after.copy_(prefix_lens + int(self.spec_num_tokens)) + + block_ids = self.block_ids_buf[:bs] + block_ids.fill_(int(self.mask_token_id)) + block_ids[:, 0].copy_(current_tokens.to(torch.int32)) + block_positions = self.block_positions_buf[:bs] + block_positions.copy_( + prefix_lens.to(torch.int64).unsqueeze(1) + self.block_offsets + ) + + cache_locs = self.draft_out_cache_loc_buf[: bs * self.spec_num_tokens] + compute_out_cache_loc_uniform( + out_cache_loc_ptr=cache_locs, + req_pool_indices=req_pool_indices, + uniform_input_length=self.spec_num_tokens, + cache_start=prefix_lens, + req_to_pages=self.req_to_page, + page_size=self.page_size, + ) + + if not (torch.cuda.is_available() and torch.cuda.is_current_stream_capturing()): + self.attn_backend.init_forward_metadata( + bs=bs, + num_extends=0, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens_after, + req_to_page=self.req_to_page, + forward_mode=ForwardMode.DECODE, + # Draft block runs in DECODE mode; the extend_* params are + # required by the signature but unused on the decode path. + extend_seq_lens=None, + extend_seq_lens_cpu=self.draft_extend_seq_lens_cpu[:bs], + extend_prefix_lens=None, + extend_prefix_lens_cpu=None, + ) + else: + # CUDA-graph capture/replay: the expanded decode metadata + # (page table) is prepared out-of-graph by the wrapper; broadcast + # the live per-request block-end length into the expanded seq_lens + # buffer here so the recorded op re-derives them on every replay. + self.attn_backend.fill_block_decode_seq_lens(bs, seq_lens_after) + + ctx = ForwardContext( + attn_backend=self.attn_backend, + token_to_kv_pool=self.token_to_kv_pool, + req_to_page=self.req_to_page, + bs=bs, + num_extends=0, + input_num_tokens=bs * self.spec_num_tokens, + forward_mode=ForwardMode.DECODE, + capture_hidden_mode=CaptureHiddenMode.FULL, + ) + + flat_ids = block_ids.reshape(-1) + input_embeds = self.embed_tokens(flat_ids) + with torch.inference_mode(): + logits_output = self.draft_model_runner.forward( + ctx=ctx, + input_ids=flat_ids, + positions=block_positions.reshape(-1), + out_cache_loc=cache_locs, + captured_hidden_states=None, + input_embeds=input_embeds, + ) + + draft_hidden = logits_output.hidden_states + if draft_hidden is None: + raise RuntimeError( + "Native DFLASH draft model did not return hidden states." + ) + draft_hidden = draft_hidden.view(bs, self.spec_num_tokens, self.hidden_size) + + next_tokens = torch.empty( + (bs, self.spec_num_tokens), dtype=torch.int32, device=self.device + ) + next_tokens[:, 0] = current_tokens.to(torch.int32) + sampled = self._greedy_sample_from_vocab_parallel_head( + draft_hidden[:, 1:, :].reshape(-1, self.hidden_size) + ) + next_tokens[:, 1:] = sampled.view(bs, self.spec_num_tokens - 1) + # Defense-in-depth: keep draft ids non-negative before they are written + # to future_input_map and embedded by the next verify forward, mirroring + # the EAGLE drafter's draft_ids.clamp_(min=0) guard. A negative id (the + # -1 NaN sentinel) would otherwise index the embedding table out of + # bounds (CUDA illegal memory access). + next_tokens.clamp_(min=0) + return next_tokens + + @nvtx_range("drafter:dflash", color="purple") + def run( + self, + base_ctx: ForwardContext, + logits_output: LogitsProcessorOutput, + output_tokens: torch.Tensor, + accept_lengths: torch.Tensor, + ) -> torch.Tensor: + if not hasattr(self, "target_model"): + raise RuntimeError("DFLASH drafter is not bound to a target model.") + self._update_native_cache_from_target(base_ctx, logits_output, accept_lengths) + current_tokens = self._current_tokens_from_output( + output_tokens, + accept_lengths, + base_ctx.num_extends, + self.spec_num_tokens, + ) + return self.draft(current_tokens) diff --git a/python/tokenspeed/runtime/execution/drafter/eagle.py b/python/tokenspeed/runtime/execution/drafter/eagle.py new file mode 100644 index 0000000..17d1422 --- /dev/null +++ b/python/tokenspeed/runtime/execution/drafter/eagle.py @@ -0,0 +1,544 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import torch +from tokenspeed_kernel.ops.sampling import argmax as sampling_argmax +from typing_extensions import override + +from tokenspeed.runtime.execution.cache_loc_kernel import ( + compute_out_cache_loc_uniform, +) +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.execution.drafter.base import BaseDrafter +from tokenspeed.runtime.execution.forward_batch_info import ( + CaptureHiddenMode, + ForwardMode, +) +from tokenspeed.runtime.multimodal.inputs import maybe_substitute_mm_pad +from tokenspeed.runtime.utils.nvtx import nvtx_range + +DsaTopKState = tuple[Any | None, Any | None] + +if TYPE_CHECKING: + from tokenspeed.runtime.execution.input_buffer import InputBuffers + from tokenspeed.runtime.execution.model_runner import ModelRunner + from tokenspeed.runtime.execution.runtime_states import RuntimeStates + from tokenspeed.runtime.layers.attention.backends.base import AttentionBackend + from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool + from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput + + +def _advance_draft_forward_metadata_if_supported(attn_backend, seq_lens) -> None: + advance = getattr(attn_backend, "advance_draft_forward_metadata", None) + if advance is not None: + advance(seq_lens) + + +@dataclass +class EagleDraftInput: + input_num_tokens: int + num_extends: int + forward_mode: ForwardMode + base_model_output: torch.Tensor # [bs] + accept_lengths: torch.Tensor # [bs] + base_out_hidden_states: torch.Tensor + global_num_tokens: list[int] | None = None + global_bs: list[int] | None = None + all_decode_or_idle: bool = False + dsa_topk: DsaTopKState = (None, None) + + +class Eagle(BaseDrafter): + """ + Draft model runner that implements the Eagle/Eagle3 algorithm. + """ + + def __init__( + self, + spec_num_tokens: int, + spec_num_steps: int, + page_size: int, + draft_model_runner: ModelRunner, + req_to_page: torch.Tensor, + attn_backend: AttentionBackend | None = None, + token_to_kv_pool: BaseTokenToKVPool | None = None, + runtime_states: RuntimeStates | None = None, + input_buffers: InputBuffers | None = None, + vocab_size: int | None = None, + ) -> None: + + super().__init__( + spec_num_tokens, + spec_num_steps, + draft_model_runner, + runtime_states=runtime_states, + input_buffers=input_buffers, + page_size=page_size, + req_to_page=req_to_page, + attn_backend=attn_backend, + token_to_kv_pool=token_to_kv_pool, + vocab_size=vocab_size, + ) + + self.device = draft_model_runner.device + hot_token_ids = draft_model_runner.model.get_hot_token_id() + + if hot_token_ids is not None: + self.hot_token_ids = hot_token_ids.to(self.device) + else: + self.hot_token_ids = None + + # For constructing fallback global_num_tokens during CUDA graph capture. + self.dp_size = draft_model_runner.mapping.attn.dp_size + self.world_size = draft_model_runner.mapping.world_size + + # Drafter-owned alias source for the draft attn backend; advanced in + # place during multi-step decode. + self.draft_seq_lens_buf = torch.zeros_like(self.input_buffers.seq_lens_buf) + + # Persistent output buffer for the draft step's compute_out_cache_loc. + self.draft_out_cache_loc_buf = torch.empty( + (self.input_buffers.max_bs * (spec_num_steps - 1),), + dtype=torch.int32, + device=self.device, + ) + + # Precomputed `arange(max_bs) * spec_num_tokens - 1` + # gather_ids = gather_ids_offsets + accept_lengths + self.padded_gather_ids_offsets_buf = ( + torch.arange( + self.input_buffers.max_bs, dtype=torch.int64, device=self.device + ) + * spec_num_tokens + - 1 + ) + + # VLM placeholder id plumbed by ModelExecutor; None for text-only targets. + self.mm_pad_substitute_id: int | None = None + hf_config = getattr(draft_model_runner.model_config, "hf_config", None) + self._dsa_reuse_mtp_topk = bool( + getattr(hf_config, "index_share_for_mtp_iteration", False) + ) + + def _accepted_output_indices( + self, + accept_lengths: torch.Tensor, + row_count: int, + *, + base_offset: int = 0, + ) -> torch.Tensor: + """Return safe flat output-token indices for each decode request. + + ``accept_lengths`` is the number of tokens that may be committed. When + the context-length cap reduces a row to 0 there is no real newly + committed output token, but the drafter still runs to preserve graph + shape. Use the row's first verify output as a valid dummy source rather + than producing ``row * N - 1``. + """ + safe_accept_lengths = ( + accept_lengths[:row_count].to(torch.int64).clamp(1, self.spec_num_tokens) + ) + return ( + self.padded_gather_ids_offsets_buf[:row_count] + + safe_accept_lengths + + base_offset + ) + + def set_mm_pad_substitute_id(self, token_id: int) -> None: + self.mm_pad_substitute_id = token_id + + # ------------------------------------------------------------------ + # Internal helpers + # ------------------------------------------------------------------ + + def _attach_dsa_topk( + self, + ctx: ForwardContext, + dsa_topk: DsaTopKState, + ) -> None: + if not self._dsa_reuse_mtp_topk: + return + ctx.dsa_prefill_topk, ctx.dsa_decode_topk = dsa_topk + + def _extract_dsa_topk( + self, + ctx: ForwardContext, + dsa_topk: DsaTopKState, + ) -> DsaTopKState: + if not self._dsa_reuse_mtp_topk: + return dsa_topk + return ctx.dsa_prefill_topk, ctx.dsa_decode_topk + + def _map_hot(self, ids: torch.Tensor) -> torch.Tensor: + """Map token ids through hot_token_ids if available, otherwise return as-is.""" + return self.hot_token_ids[ids] if self.hot_token_ids is not None else ids + + def _get_first_step_input( + self, + draft_input: EagleDraftInput, + bs: int, + input_num_tokens: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Returns (input_ids, gather_ids) for the first draft step. + + The first-step input shape matches the base model's: ragged + ``[prefill_part || decode_part]`` under MIXED, full prefill chunks + under EXTEND, ``base_model_output`` directly under DECODE. + """ + num_extends = draft_input.num_extends + num_decodes = bs - num_extends + if num_extends > 0: + num_decode_tokens = num_decodes * self.spec_num_tokens + num_prefill_tokens = input_num_tokens - num_decode_tokens + + input_ids = self.input_buffers.shifted_prefill_ids_buf[:input_num_tokens] + unpadded_input_lengths = self.input_buffers.input_lengths_buf[:bs] + if num_decodes > 0: + input_ids[num_prefill_tokens:].copy_( + draft_input.base_model_output[num_extends:] + ) + unpadded_input_lengths[num_extends:].copy_( + draft_input.accept_lengths[num_extends:] + ) + + last_indices = unpadded_input_lengths[:num_extends].cumsum(0) - 1 + last_input_ids = input_ids[last_indices] + input_ids[last_indices] = torch.where( + last_input_ids == -1, + draft_input.base_model_output[:num_extends], + last_input_ids, + ) + + gather_ids = last_indices + if num_decodes > 0: + gather_ids = torch.cat( + [ + gather_ids, + self._accepted_output_indices( + draft_input.accept_lengths[num_extends:], + num_decodes, + base_offset=num_prefill_tokens, + ), + ] + ) + else: + input_ids = draft_input.base_model_output + gather_ids = self._accepted_output_indices( + draft_input.accept_lengths, + bs, + ) + + return input_ids, gather_ids + + @nvtx_range("draft_first_step", color="purple") + def _run_first_step( + self, + bs: int, + draft_input: EagleDraftInput, + ) -> tuple[LogitsProcessorOutput, DsaTopKState]: + + buffers = self.input_buffers + forward_mode = draft_input.forward_mode + + input_ids, gather_ids = self._get_first_step_input( + draft_input, bs, draft_input.input_num_tokens + ) + input_ids = maybe_substitute_mm_pad(input_ids, self.mm_pad_substitute_id) + draft_model = self.draft_model_runner.model + input_num_tokens = draft_input.input_num_tokens + + ctx = ForwardContext( + attn_backend=self.attn_backend, + token_to_kv_pool=self.token_to_kv_pool, + req_to_page=self.req_to_page, + bs=bs, + num_extends=draft_input.num_extends, + input_num_tokens=input_num_tokens, + forward_mode=forward_mode, + capture_hidden_mode=CaptureHiddenMode.LAST, + gather_ids=gather_ids, + global_num_tokens=draft_input.global_num_tokens, + global_bs=draft_input.global_bs, + all_decode_or_idle=draft_input.all_decode_or_idle, + draft_seq_lens_buf=self.draft_seq_lens_buf, + accept_lengths=draft_input.accept_lengths, + ) + + dsa_topk = draft_input.dsa_topk + prepare_dsa_topk = getattr(draft_model, "prepare_dsa_topk_for_mtp_decode", None) + compute_dsa_topk_first_step = bool( + getattr(draft_model, "compute_dsa_topk_first_step", False) + ) + if compute_dsa_topk_first_step: + # GLM NextN has its own indexer weights. Compute first-step top-k + # in the draft model, then select rows used by later MTP steps. + dsa_topk = (None, None) + elif draft_input.num_extends == 0 and prepare_dsa_topk is not None: + dsa_topk = prepare_dsa_topk(dsa_topk, gather_ids) + else: + dsa_topk = (None, None) + self._attach_dsa_topk(ctx, dsa_topk) + + logits_output = self.draft_model_runner.forward( + ctx=ctx, + input_ids=input_ids, + positions=buffers.positions_buf[:input_num_tokens], + out_cache_loc=buffers.out_cache_loc_buf[:input_num_tokens], + captured_hidden_states=draft_input.base_out_hidden_states, + spec_step_idx=0, + ) + dsa_topk = self._extract_dsa_topk(ctx, dsa_topk) + if compute_dsa_topk_first_step and prepare_dsa_topk is not None: + dsa_topk = prepare_dsa_topk( + dsa_topk, + gather_ids, + num_prefill_rows=draft_input.num_extends, + ) + return logits_output, dsa_topk + + @nvtx_range("draft_multi_step", color="purple") + def _run_multi_step_decode( + self, + bs: int, + draft_ids: torch.Tensor, + next_tokens: torch.Tensor, + logits_output: LogitsProcessorOutput, + draft_input: EagleDraftInput, + dsa_topk: DsaTopKState, + ) -> None: + num_extends = draft_input.num_extends + num_decodes = bs - num_extends + req_pool_indices = self.input_buffers.req_pool_indices_buf[:bs] + cache_start = self.input_buffers.seq_lens_buf[:bs].clone() + # Step 1's write position uses vc+accept_length after target verify so + # rotary/cache metadata stay on the accepted prefix, not rejected tail. + if num_decodes > 0: + cache_start[num_extends:] = ( + self.runtime_states.valid_cache_lengths.index_select( + 0, req_pool_indices[num_extends:] + ) + + draft_input.accept_lengths[num_extends:] + ) + + # Write cache slots for steps 1..N-1. + cache_locs = self.draft_out_cache_loc_buf[: bs * (self.spec_num_steps - 1)] + compute_out_cache_loc_uniform( + out_cache_loc_ptr=cache_locs, + req_pool_indices=req_pool_indices, + uniform_input_length=self.spec_num_steps - 1, + cache_start=cache_start, + req_to_pages=self.req_to_page, + page_size=self.page_size, + ) + cache_locs = cache_locs.view(bs, self.spec_num_steps - 1) + # +1 is the kernel's read-inclusive convention; advanced per iter. + draft_seq_lens = self.draft_seq_lens_buf[:bs] + torch.add(cache_start, 1, out=draft_seq_lens) + + positions = cache_start.clone() + + for i in range(1, self.spec_num_steps): + # make a ctx every time model runner forward + # Multi-step decode is pure DECODE mode: one token per request. + # global_num_tokens must reflect each rank's batch size, not the + # target model's total tokens (which may be bs * spec_num_tokens). + global_num_tokens = draft_input.global_num_tokens + + if self.dp_size > 1: + if draft_input.global_bs is not None: + global_num_tokens = draft_input.global_bs + else: + # CUDA graph capture path: uniform batch size across ranks. + global_num_tokens = [bs] * self.world_size + + ctx = ForwardContext( + bs=bs, + num_extends=0, + attn_backend=self.attn_backend, + token_to_kv_pool=self.token_to_kv_pool, + req_to_page=self.req_to_page, + input_num_tokens=bs, + forward_mode=ForwardMode.DECODE, + capture_hidden_mode=CaptureHiddenMode.LAST, + global_num_tokens=global_num_tokens, + global_bs=draft_input.global_bs, + all_decode_or_idle=draft_input.all_decode_or_idle, + ) + self._attach_dsa_topk(ctx, dsa_topk) + + out_cache_loc = cache_locs[:, i - 1].contiguous() + # Keep attention metadata on the accepted prefix; rejected verify + # tail slots may still contain stale draft KV. + _advance_draft_forward_metadata_if_supported( + ctx.attn_backend, + draft_seq_lens, + ) + + with nvtx_range("draft_forward", color="red"): + logits_output = self.draft_model_runner.forward( + ctx=ctx, + input_ids=self._map_hot(draft_ids), + positions=positions, + out_cache_loc=out_cache_loc, + captured_hidden_states=logits_output.hidden_states, + spec_step_idx=i, + ) + dsa_topk = self._extract_dsa_topk(ctx, dsa_topk) + + with nvtx_range("draft_sample", color="yellow"): + if logits_output.next_token_ids is not None: + draft_ids = logits_output.next_token_ids + else: + draft_ids = sampling_argmax(logits_output.next_token_logits) + # Column 0 holds last_verified_ids; drafter writes step `i` into column `i + 1`. + next_tokens[:, i + 1] = self._map_hot(draft_ids) + if i + 1 < self.spec_num_steps: + positions.add_(1) + draft_seq_lens.add_(1) + + # ------------------------------------------------------------------ + # Public entry point (type-based dispatch from ModelExecutor) + # ------------------------------------------------------------------ + + @override + def get_candidates( + self, + base_ctx: ForwardContext, + ) -> torch.Tensor | None: + num_extends = base_ctx.num_extends + num_decodes = base_ctx.bs - num_extends + if num_decodes == 0: + return None + + num_decode_tokens = num_decodes * self.spec_num_tokens + num_prefill_tokens = base_ctx.input_num_tokens - num_decode_tokens + return self.input_buffers.input_ids_buf[ + num_prefill_tokens : base_ctx.input_num_tokens + ].reshape(num_decodes, self.spec_num_tokens) + + @override + def draft( + self, + draft_input: EagleDraftInput, + ) -> torch.Tensor: + + bs = draft_input.accept_lengths.shape[0] + + # Layout: column 0 holds the last verified id (the base model's accepted token); + # columns 1..spec_num_steps hold the drafter's speculative tokens. + next_tokens = torch.empty( + (bs, self.spec_num_steps + 1), + dtype=torch.int32, + device=self.device, + ) + + # Last verified id per request → next_tokens[:, 0]. + num_extends = draft_input.num_extends + num_decodes = bs - num_extends + if num_extends > 0: + next_tokens[:num_extends, 0] = draft_input.base_model_output[:num_extends] + if num_decodes > 0: + indices = self._accepted_output_indices( + draft_input.accept_lengths[num_extends:], + num_decodes, + ) + if num_extends > 0: + indices.add_(num_extends) + torch.index_select( + draft_input.base_model_output, + 0, + indices, + out=next_tokens[num_extends:, 0], + ) + if self.spec_num_steps > 0: + next_tokens[:, 1:] = next_tokens[:, :1] + + # Seed the draft attn backend's aliased seq_lens for the first step. + self.draft_seq_lens_buf[:bs].copy_(self.input_buffers.seq_lens_buf[:bs]) + + # First draft step. LogitsProcessor prunes `[num_prefill_tokens + num_decodes * spec_num_tokens, ...]` + # down to `[bs, ...]`, so logits/hidden_states arrive here already aligned to one row per request. + logits_output, dsa_topk = self._run_first_step(bs, draft_input) + + if logits_output.next_token_ids is not None: + draft_ids = logits_output.next_token_ids + else: + draft_ids = sampling_argmax(logits_output.next_token_logits) + next_tokens[:, 1] = self._map_hot(draft_ids) + + if self.spec_num_steps <= 1: + return next_tokens + + if self.input_buffers.all_extends_mid_chunk and self.dp_size == 1: + # Skip multi-step when the whole batch is mid-chunk EXTEND: + # no request completes a target-side speculative verification + # after this forward, so any speculative tokens would be discarded. + # + # In DP we still run, because peer ranks may have completing + # extends or decodes; diverging here would desync the drafter's + # dense-TP / MoE-EP collectives (NCCL hang or RSAG mismatch). + return next_tokens + + # Draft step 2+ (multi-step decode). + # Multi-step decode operates on full bs; drop the [num_extends:] + # slice that step 0 may have set up for MIXED target. No-op on + # backends that fill separate prefill/decode metadata at init + # time. + with self.attn_backend.override_num_extends(0): + self._run_multi_step_decode( + bs, + draft_ids, + next_tokens, + logits_output, + draft_input, + dsa_topk, + ) + return next_tokens + + @override + @nvtx_range("drafter", color="purple") + def run( + self, + base_ctx: ForwardContext, + logits_output: LogitsProcessorOutput, + output_tokens: torch.Tensor, + accept_lengths: torch.Tensor, + ) -> torch.Tensor: + + draft_input = EagleDraftInput( + input_num_tokens=base_ctx.input_num_tokens, + num_extends=base_ctx.num_extends, + forward_mode=base_ctx.forward_mode, + base_model_output=output_tokens, + accept_lengths=accept_lengths, + base_out_hidden_states=logits_output.hidden_states, + global_num_tokens=base_ctx.global_num_tokens, + global_bs=base_ctx.global_bs, + all_decode_or_idle=base_ctx.all_decode_or_idle, + dsa_topk=(base_ctx.dsa_prefill_topk, base_ctx.dsa_decode_topk), + ) + + # next_tokens layout: column 0 = last verified id, columns 1.. = drafter tokens. + return self.draft(draft_input) diff --git a/python/tokenspeed/runtime/execution/factory.py b/python/tokenspeed/runtime/execution/factory.py new file mode 100644 index 0000000..0b457ca --- /dev/null +++ b/python/tokenspeed/runtime/execution/factory.py @@ -0,0 +1,113 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Factory helpers for model runners and model executors.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import tokenspeed.runtime.layers.attention.backends # noqa: F401 # trigger register_backend() calls +from tokenspeed.runtime.configs.model_config import ModelConfig +from tokenspeed.runtime.execution.model_executor import ( + ModelExecutor, + ModelExecutorConfig, +) +from tokenspeed.runtime.execution.model_runner import ModelRunner +from tokenspeed.runtime.sampling.registry import create_sampling_backend +from tokenspeed.runtime.utils.nvtx import set_nvtx_enabled +from tokenspeed.runtime.utils.server_args import ServerArgs + +if TYPE_CHECKING: + from tokenspeed.runtime.layers.attention.backends.base import AttentionBackend + from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool + + +def create_model_runner( + server_args: ServerArgs, + model_config: ModelConfig, + draft_model_config: ModelConfig | None, + gpu_id: int, + global_rank: int, +): + """Create the main model runner and optional draft model runner.""" + model_runner = ModelRunner( + model_config=model_config, + gpu_id=gpu_id, + server_args=server_args, + global_rank=global_rank, + ) + + draft_model_runner = None + if draft_model_config is not None: + draft_model_runner = ModelRunner( + model_config=draft_model_config, + gpu_id=gpu_id, + server_args=server_args, + global_rank=global_rank, + is_draft_worker=True, + ) + + return model_runner, draft_model_runner + + +def create_model_executor( + server_args: ServerArgs, + config: ModelExecutorConfig, + model_runner: ModelRunner, + attn_backend: AttentionBackend, + token_to_kv_pool: BaseTokenToKVPool, + draft_model_runner: ModelRunner | None = None, + draft_attn_backend: AttentionBackend | None = None, + draft_token_to_kv_pool: BaseTokenToKVPool | None = None, + mamba_pool: object | None = None, +) -> ModelExecutor: + """Create the model executor with its sampler configuration.""" + if server_args.enable_nvtx: + set_nvtx_enabled(True) + + max_bs = config.max_num_seqs // max(config.data_parallel_size, 1) + + max_draft_tokens_per_req = ( + config.spec_num_tokens if config.spec_algo is not None else 1 + ) + + sampling_backend = create_sampling_backend( + server_args, + max_bs=max_bs, + max_draft_tokens_per_req=max_draft_tokens_per_req, + device=config.device, + max_req_pool_size=config.max_req_pool_size, + vocab_size=config.vocab_size, + # Same TP group as LogitsProcessor. + tp_group=model_runner.mapping.attn.tp_group, + ) + + return ModelExecutor( + config=config, + model_runner=model_runner, + attn_backend=attn_backend, + token_to_kv_pool=token_to_kv_pool, + sampling_backend=sampling_backend, + draft_model_runner=draft_model_runner, + draft_attn_backend=draft_attn_backend, + draft_token_to_kv_pool=draft_token_to_kv_pool, + mamba_pool=mamba_pool, + ) diff --git a/python/tokenspeed/runtime/execution/forward_batch_info.py b/python/tokenspeed/runtime/execution/forward_batch_info.py new file mode 100755 index 0000000..4539789 --- /dev/null +++ b/python/tokenspeed/runtime/execution/forward_batch_info.py @@ -0,0 +1,145 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Forward mode enums and position computation helpers.""" + +from __future__ import annotations + +from enum import IntEnum, auto + +import torch +import triton +import triton.language as tl + + +class ForwardMode(IntEnum): + # Extend a sequence. The KV cache of the beginning part of the sequence + # is already computed (e.g., system prompt). + EXTEND = auto() + # Decode one or more tokens per request. + DECODE = auto() + # Contains both EXTEND and DECODE tokens in one batch. + MIXED = auto() + # No sequence to forward; used for data parallel attention idle ranks. + IDLE = auto() + + def is_extend(self): + return self == ForwardMode.EXTEND + + def is_decode(self): + return self == ForwardMode.DECODE + + def is_mixed(self): + return self == ForwardMode.MIXED + + def is_idle(self): + return self == ForwardMode.IDLE + + def is_extend_or_mixed(self): + return self == ForwardMode.EXTEND or self == ForwardMode.MIXED + + def is_decode_or_idle(self): + return self == ForwardMode.DECODE or self == ForwardMode.IDLE + + @staticmethod + def from_num_extends(num_extends: int, batch_size: int) -> "ForwardMode": + if batch_size <= 0: + return ForwardMode.IDLE + elif num_extends > 0: + return ForwardMode.MIXED if num_extends < batch_size else ForwardMode.EXTEND + else: + return ForwardMode.DECODE + + +class CaptureHiddenMode(IntEnum): + NULL = auto() + # Capture hidden states of all tokens. + FULL = auto() + # Capture a hidden state of the last token. + LAST = auto() + + def need_capture(self): + return self != CaptureHiddenMode.NULL + + def is_full(self): + return self == CaptureHiddenMode.FULL + + def is_last(self): + return self == CaptureHiddenMode.LAST + + +def compute_position_triton( + extend_prefix_lens: torch.Tensor, + extend_seq_lens: torch.Tensor, + extend_seq_lens_sum, + out: torch.Tensor | None = None, +): + batch_size = extend_seq_lens.shape[0] + if out is None: + positions = torch.empty( + extend_seq_lens_sum, dtype=torch.int64, device=extend_seq_lens.device + ) + else: + if out.numel() < extend_seq_lens_sum: + raise ValueError( + "compute_position_triton out buffer too small: " + f"{out.numel()} < {extend_seq_lens_sum}" + ) + positions = out + extend_start_loc = torch.empty( + batch_size, dtype=torch.int32, device=extend_seq_lens.device + ) + has_prefix = extend_prefix_lens.shape[0] == batch_size + # Launch kernel + compute_position_kernel[(batch_size,)]( + positions, extend_start_loc, extend_prefix_lens, extend_seq_lens, has_prefix + ) + + return positions, extend_start_loc + + +@triton.jit +def compute_position_kernel( + positions, + extend_start_loc, + extend_prefix_lens, + extend_seq_lens, + has_prefix: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 512 + pid = tl.program_id(0).to(tl.int64) + + prefix_len = tl.load(extend_prefix_lens + pid) if has_prefix else 0 + seq_len = tl.load(extend_seq_lens + pid) + + # This can be slow for large bs + cumsum_start = tl.cast(0, tl.int64) + for i in range(pid): + cumsum_start += tl.load(extend_seq_lens + i) + + num_loop = tl.cdiv(seq_len, BLOCK_SIZE) + for i in range(num_loop): + offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE + tl.store( + positions + cumsum_start + offset, + prefix_len + offset, + mask=offset < seq_len, + ) + tl.store(extend_start_loc + pid, cumsum_start) diff --git a/python/tokenspeed/runtime/execution/input_buffer.py b/python/tokenspeed/runtime/execution/input_buffer.py new file mode 100644 index 0000000..9f7bd96 --- /dev/null +++ b/python/tokenspeed/runtime/execution/input_buffer.py @@ -0,0 +1,453 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from tokenspeed.runtime.execution.cache_loc_kernel import ( + compute_out_cache_loc, + fused_decode_input_prep, +) +from tokenspeed.runtime.execution.forward_batch_info import compute_position_triton +from tokenspeed.runtime.utils import get_colorful_logger +from tokenspeed.runtime.utils.nvtx import nvtx_range + +if TYPE_CHECKING: + from tokenspeed.runtime.execution.runtime_states import RuntimeStates + + +logger = get_colorful_logger(__name__) + + +class InputBuffers: + """ + ForwardContext tensor data source, read-only after fill. Holds only + model-forward inputs; per-request sampling scalars (temperature, top_k, + penalties, seed, etc.) live on the sampling backend as pool-indexed + buffers populated on slot flips. + """ + + def __init__( + self, + max_bs: int, + max_num_tokens: int, + page_size: int, + dummy_kv_slot: int, + state_write_padding_pool_index: int, + device: str = "cuda", + has_mamba: bool = False, + ): + self.device = device + self.page_size = page_size + self.max_num_tokens = max_num_tokens + self.dummy_kv_slot = dummy_kv_slot + self.state_write_padding_pool_index = state_write_padding_pool_index + self.max_bs = max_bs + self.all_extends_mid_chunk = False + self.has_mamba = has_mamba + + with torch.device(device): + # Initialise buffers to the *padding* values the captured graph + # expects for padded rows (input_ids=1, positions=0, req_pool=0, + # seq_lens=1, out_cache_loc=dummy_kv_slot). Each iteration overwrites + # the active prefix [:total_tokens]; fill_input_buffers refreshes the + # padding tail [total_tokens:] back to these defaults every step, + # because a larger prior iter can leave stale values past the + # current prefix. + self.input_ids_buf = torch.ones((max_num_tokens,), dtype=torch.int32) + # Used in draft prefill + self.shifted_prefill_ids_buf = torch.ones_like(self.input_ids_buf) + self.input_lengths_buf = torch.ones((max_num_tokens,), dtype=torch.int32) + # Zero (not arange) so padded positions read a consistent, in-range + # value; the tail is re-zeroed every iteration by fill_input_buffers. + self.positions_buf = torch.zeros(max_num_tokens, dtype=torch.int64) + self.mrope_positions_buf = torch.zeros( + (3, max_num_tokens), dtype=torch.int64 + ) + self.req_pool_indices_buf = torch.zeros((max_bs,), dtype=torch.int64) + self.state_write_req_pool_indices_buf = torch.full( + (max_bs,), state_write_padding_pool_index, dtype=torch.int64 + ) + self.seq_lens_buf = torch.ones((max_bs,), dtype=torch.int32) + # Initialise to dummy_kv_slot so that padding positions (never + # written by compute_out_cache_loc) always point to the reserved + # dummy KV slot and never corrupt real KV cache entries. + self.out_cache_loc_buf = torch.full( + (max_num_tokens,), dummy_kv_slot, dtype=torch.int32 + ) + self.force_single_token_verify_buf = torch.zeros(max_bs, dtype=torch.bool) + self.extend_prefix_lens_buf = torch.zeros(max_bs, dtype=torch.int32) + self.extend_seq_lens_buf = torch.zeros(max_bs, dtype=torch.int32) + if has_mamba: + self.mamba_pool_indices_buf = torch.full( + (max_bs,), -1, dtype=torch.int32 + ) + self.mamba_cow_src_indices_buf = torch.full( + (max_bs,), -1, dtype=torch.int32 + ) + self.mamba_branching_seqlens_buf = torch.full( + (max_bs,), -1, dtype=torch.int32 + ) + self.mamba_track_pool_indices_buf = torch.full( + (max_bs,), -1, dtype=torch.int32 + ) + + self.extend_prefix_lens_cpu = torch.zeros( + max_bs, dtype=torch.int32, pin_memory=True + ) + self.extend_seq_lens_cpu = torch.zeros( + max_bs, dtype=torch.int32, pin_memory=True + ) + if has_mamba: + self._mamba_pool_indices_cpu = torch.full( + (max_bs,), -1, dtype=torch.int32, pin_memory=True + ) + self._mamba_cow_src_indices_cpu = torch.full( + (max_bs,), -1, dtype=torch.int32, pin_memory=True + ) + self._mamba_branching_seqlens_cpu = torch.full( + (max_bs,), -1, dtype=torch.int32, pin_memory=True + ) + self._mamba_track_pool_indices_cpu = torch.full( + (max_bs,), -1, dtype=torch.int32, pin_memory=True + ) + + @nvtx_range("input_prep_fill", color="cyan") + def fill_input_buffers( + self, + forward_op, + runtime_states: RuntimeStates, + req_to_page: torch.Tensor, + total_tokens: int, + ): + batch_size = len(forward_op.request_ids) + num_extends = forward_op.num_extends() + + # CPU-side fast path: when the scheduler always emits a decode_input_ids + # list (even though every entry is -1, meaning "no override"). + decode_input_ids = forward_op.decode_input_ids + if decode_input_ids is not None and all(x == -1 for x in decode_input_ids): + decode_input_ids = None + req_pool_indices_cpu = torch.tensor( + forward_op.request_pool_indices, device="cpu", pin_memory=True + ) + self.req_pool_indices_buf[:batch_size].copy_( + req_pool_indices_cpu, + non_blocking=True, + ) + self.state_write_req_pool_indices_buf[:batch_size].copy_( + req_pool_indices_cpu, + non_blocking=True, + ) + input_lengths_cpu = torch.tensor( + forward_op.input_lengths, + dtype=torch.int32, + device="cpu", + pin_memory=True, + ) + self.input_lengths_buf[:batch_size].copy_( + input_lengths_cpu, + non_blocking=True, + ) + + self.all_extends_mid_chunk = ( + num_extends > 0 + and num_extends == batch_size + and all( + forward_op.extend_prefix_lens[i] + forward_op.input_lengths[i] + < forward_op.prefill_lengths[i] + for i in range(num_extends) + ) + ) + + if num_extends > 0: + self.extend_prefix_lens_cpu[:num_extends] = torch.as_tensor( + forward_op.extend_prefix_lens, dtype=torch.int32 + ) + self.extend_prefix_lens_buf[:num_extends].copy_( + self.extend_prefix_lens_cpu[:num_extends], non_blocking=True + ) + self.extend_seq_lens_cpu[:num_extends] = torch.as_tensor( + forward_op.input_lengths[:num_extends], dtype=torch.int32 + ) + self.extend_seq_lens_buf[:num_extends].copy_( + self.extend_seq_lens_cpu[:num_extends], non_blocking=True + ) + + # Get valid cache lengths for requests + req_pool_indices_device = self.req_pool_indices_buf[:batch_size] + input_lengths_device = self.input_lengths_buf[:batch_size] + + def write_decode_input_ids( + decode_req_pool_indices: torch.Tensor, + decode_input_ids: list[int], + row_offset: int, + expected_count: int, + context: str, + ) -> None: + if len(decode_input_ids) != expected_count: + raise RuntimeError( + f"{context} decode_input_ids length mismatch: " + f"got {len(decode_input_ids)}, expected {expected_count}" + ) + decode_input_ids_tensor = torch.tensor( + decode_input_ids, + dtype=torch.int32, + device="cpu", + pin_memory=True, + ).to(req_pool_indices_device.device, non_blocking=True) + mask = (decode_input_ids_tensor != -1).unsqueeze(1) + ids = decode_input_ids_tensor.unsqueeze(1) + + # Col 0: verified token (mask preserves drafter-owned rows). + first_slot = runtime_states.future_input_map[decode_req_pool_indices, :1] + runtime_states.future_input_map[decode_req_pool_indices, :1] = torch.where( + mask, ids, first_slot + ) + # Cols 1.. are real candidates only when the local drafter or the + # remote P-side path populated them. Bootstrap/recovery rows with + # no candidate source still feed a full-width target forward, so + # use a valid dummy token in model inputs and force the verifier to + # consume only the first target token for those rows. + width = runtime_states.future_input_map.shape[1] + remote_candidate_ready = runtime_states.remote_spec_candidate_ready[ + decode_req_pool_indices + ] + force_single_token = mask.squeeze(1) & ~remote_candidate_ready + if width > 1: + tail = runtime_states.future_input_map[decode_req_pool_indices, 1:] + dummy_tail = ids.expand(-1, width - 1) + runtime_states.future_input_map[decode_req_pool_indices, 1:] = ( + torch.where(force_single_token.unsqueeze(1), dummy_tail, tail) + ) + self.force_single_token_verify_buf[ + row_offset : row_offset + expected_count + ] = force_single_token + runtime_states.remote_spec_candidate_ready[decode_req_pool_indices] = False + + # Decode-only fast path: one fused Triton kernel writes out_cache_loc, + # positions, and seq_lens in a single launch and reads + # valid_cache_lengths[pool_idx] directly, so the indexSelect + cumsum + # path + compute_position + seq_lens add are all gone. + if num_extends == 0 and batch_size > 0: + fused_decode_input_prep( + out_cache_loc_ptr=self.out_cache_loc_buf[:total_tokens], + positions_ptr=self.positions_buf[:total_tokens], + seq_lens_out_ptr=self.seq_lens_buf[:batch_size], + req_pool_indices=req_pool_indices_device, + valid_cache_lengths=runtime_states.valid_cache_lengths, + uniform_input_length=total_tokens // batch_size, + req_to_pages=req_to_page, + page_size=self.page_size, + ) + # Decode path's seq_lens / positions / out_cache_loc are done. + valid_cache_lengths = None + else: + # Mixed / pure-prefill: keep the per-kernel pipeline. indexSelect + # for valid_cache_lengths is required because compute_position and + # the seq_lens add use it. + valid_cache_lengths = runtime_states.valid_cache_lengths.index_select( + 0, req_pool_indices_device + ) + compute_out_cache_loc( + out_cache_loc_ptr=self.out_cache_loc_buf[:total_tokens], + req_pool_indices=req_pool_indices_device, + input_lengths=input_lengths_device, + cache_start=valid_cache_lengths, + req_to_pages=req_to_page, + page_size=self.page_size, + ) + + # Compute positions. In mixed batches, prefill rows use their extend + # prefix lengths while decode rows use the current valid cache lengths. + prefill_prefix_lens = self.extend_prefix_lens_buf[:num_extends] + if num_extends == batch_size: + prefix_lens = prefill_prefix_lens + else: + prefix_lens = valid_cache_lengths.clone() + prefix_lens[:num_extends].copy_(prefill_prefix_lens) + # Write positions directly into the persistent buffer to skip the + # otherwise-required DtoD copy. + compute_position_triton( + extend_prefix_lens=prefix_lens, + extend_seq_lens=input_lengths_device, + extend_seq_lens_sum=total_tokens, + out=self.positions_buf[:total_tokens], + ) + + # Determine input_ids and forward_mode + if num_extends > 0: + prefill_token_count = sum(forward_op.input_lengths[:num_extends]) + input_ids_cpu = torch.tensor( + forward_op.input_ids, device="cpu", pin_memory=True + ) + self.input_ids_buf[:prefill_token_count].copy_( + input_ids_cpu, + non_blocking=True, + ) + shifted_ids_cpu = torch.tensor( + forward_op.shifted_input_ids, device="cpu", pin_memory=True + ) + self.shifted_prefill_ids_buf[:prefill_token_count].copy_( + shifted_ids_cpu, + non_blocking=True, + ) + if num_extends < batch_size: + decode_req_pool_indices = req_pool_indices_device[ + num_extends:batch_size + ] + if decode_input_ids is not None: + write_decode_input_ids( + decode_req_pool_indices, + decode_input_ids, + num_extends, + batch_size - num_extends, + "mixed forward", + ) + decode_ids = runtime_states.future_input_map[ + decode_req_pool_indices + ].flatten() + self.input_ids_buf[prefill_token_count:total_tokens].copy_( + decode_ids, + non_blocking=True, + ) + self.shifted_prefill_ids_buf[prefill_token_count:total_tokens].copy_( + decode_ids, + non_blocking=True, + ) + else: + # If the scheduler provides explicit decode input ids (!= -1), write + # them into future_input_map before reading, so that they take effect + # as the input for this decode step. + if decode_input_ids is not None: + write_decode_input_ids( + req_pool_indices_device, + decode_input_ids, + 0, + batch_size, + "decode forward", + ) + self.input_ids_buf[:total_tokens].copy_( + runtime_states.future_input_map[req_pool_indices_device].flatten(), + non_blocking=True, + ) + + # Defensive clamp into the valid vocab range. The decode input ids come + # from future_input_map, written by the previous iteration's + # sampler/drafter; the intermittent spec-decode decode-state race can + # surface a stale/corrupt out-of-range id there. Feeding an out-of-range + # id to the captured graph's embedding gather trips a device-side assert + # (`vectorized_gather_kernel index out of bounds`) that tears the whole + # server down. Clamp the active prefix before the graph reads these + # buffers (a no-op for legitimate ids). Mirrors the post-graph + # output_tokens clamp in the output_d2h step of + # ModelExecutor.execute_forward_op. + vocab_size = runtime_states.vocab_size + self.input_ids_buf[:total_tokens].clamp_(0, vocab_size - 1) + self.shifted_prefill_ids_buf[:total_tokens].clamp_(0, vocab_size - 1) + + if valid_cache_lengths is not None: + torch.add( + input_lengths_device, + valid_cache_lengths, + out=self.seq_lens_buf[:batch_size], + ) + + # Refresh the padding tail of the persistent buffers every iteration. + # The captured graph replays at a padded batch size and DOES read the + # padded rows; a previous iter with a *larger* total_tokens / batch_size + # leaves stale values in the tail (real cache locations, per-request seq + # lengths, positions, token ids, req-pool slots). Reusing those for + # padded tokens routes KV writes into real cache slots (corruption), + # forces attention to scan oversize ranges, and -- for a stale + # out-of-range token id -- trips the embedding gather's device-side + # assert that tears the server down. The __init__ safe defaults + # (input_ids=1, req_pool=0, positions=0) are not enough on their own + # once a larger iter has overwritten the tail, so scrub it back here + # (cheap tail-only fills; the active prefix was written above). + if total_tokens < self.max_num_tokens: + self.input_ids_buf[total_tokens:].fill_(1) + self.out_cache_loc_buf[total_tokens:].fill_(self.dummy_kv_slot) + self.positions_buf[total_tokens:].fill_(0) + self.mrope_positions_buf[:, total_tokens:].zero_() + if batch_size < self.max_bs: + self.req_pool_indices_buf[batch_size:].fill_(0) + self.state_write_req_pool_indices_buf[batch_size:].fill_( + self.state_write_padding_pool_index + ) + self.seq_lens_buf[batch_size:].fill_(1) + + if ( + self.has_mamba + and hasattr(forward_op, "mamba_pool_indices") + and forward_op.mamba_pool_indices + ): + self._mamba_pool_indices_cpu[:batch_size].copy_( + torch.as_tensor(forward_op.mamba_pool_indices, dtype=torch.int32) + ) + self._mamba_cow_src_indices_cpu[:batch_size].copy_( + torch.as_tensor(forward_op.mamba_cow_src_indices, dtype=torch.int32) + ) + self._mamba_branching_seqlens_cpu[:batch_size].copy_( + torch.as_tensor(forward_op.mamba_branching_seqlens, dtype=torch.int32) + ) + self._mamba_track_pool_indices_cpu[:batch_size].copy_( + torch.as_tensor(forward_op.mamba_track_pool_indices, dtype=torch.int32) + ) + + self.mamba_pool_indices_buf[:batch_size].copy_( + self._mamba_pool_indices_cpu[:batch_size], non_blocking=True + ) + self.mamba_cow_src_indices_buf[:batch_size].copy_( + self._mamba_cow_src_indices_cpu[:batch_size], non_blocking=True + ) + self.mamba_branching_seqlens_buf[:batch_size].copy_( + self._mamba_branching_seqlens_cpu[:batch_size], non_blocking=True + ) + self.mamba_track_pool_indices_buf[:batch_size].copy_( + self._mamba_track_pool_indices_cpu[:batch_size], non_blocking=True + ) + if batch_size < self.mamba_pool_indices_buf.shape[0]: + self.mamba_pool_indices_buf[batch_size:].fill_(-1) + self.mamba_cow_src_indices_buf[batch_size:].fill_(-1) + self.mamba_branching_seqlens_buf[batch_size:].fill_(-1) + self.mamba_track_pool_indices_buf[batch_size:].fill_(-1) + + return decode_input_ids + + def fill_dummy_decode_buffers(self, batch_size: int, total_tokens: int): + """Prepare padded decode graph inputs for a rank with no real tokens.""" + if total_tokens > 0: + self.input_ids_buf[:total_tokens].fill_(1) + self.out_cache_loc_buf[:total_tokens].fill_(self.dummy_kv_slot) + self.positions_buf[:total_tokens].fill_(0) + self.mrope_positions_buf[:, :total_tokens].zero_() + if batch_size > 0: + self.req_pool_indices_buf[:batch_size].fill_(0) + self.state_write_req_pool_indices_buf[:batch_size].fill_( + self.state_write_padding_pool_index + ) + # seq_lens must be >= spec_num_tokens so the drafter's prewrite + # correction never goes negative. + num_tokens_per_req = total_tokens // batch_size if batch_size > 0 else 1 + self.seq_lens_buf[:batch_size].fill_(max(num_tokens_per_req, 1)) diff --git a/python/tokenspeed/runtime/execution/model_executor.py b/python/tokenspeed/runtime/execution/model_executor.py new file mode 100644 index 0000000..ff528c8 --- /dev/null +++ b/python/tokenspeed/runtime/execution/model_executor.py @@ -0,0 +1,2198 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import time +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch +from tokenspeed_kernel.ops.tuning import freeze_autotuning +from tokenspeed_kernel.platform import current_platform + +from tokenspeed.runtime.configs.model_config import ModelConfig +from tokenspeed.runtime.configs.paged_cache_spec import ( + scheduler_ext_flat_kvcache, + validate_flat_scheduler_config, +) +from tokenspeed.runtime.configs.utils import get_rope_parameters +from tokenspeed.runtime.engine.scheduler_utils import ( + flat_block_tables_from_forward_op, + paged_cache_block_table_base_offsets_from_forward_op, + paged_cache_block_tables_from_forward_op, +) +from tokenspeed.runtime.execution.cache_loc_kernel import update_block_table +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.execution.cuda_graph_wrapper import CudaGraphWrapper +from tokenspeed.runtime.execution.drafter.dflash import DFlash +from tokenspeed.runtime.execution.drafter.eagle import Eagle +from tokenspeed.runtime.execution.forward_batch_info import ( + CaptureHiddenMode, + ForwardMode, +) +from tokenspeed.runtime.execution.input_buffer import InputBuffers +from tokenspeed.runtime.execution.model_runner import ModelRunner +from tokenspeed.runtime.execution.nan_guard import NanGuard +from tokenspeed.runtime.execution.prefill_graph import PrefillGraph +from tokenspeed.runtime.execution.runtime_states import RuntimeStates +from tokenspeed.runtime.execution.types import ModelExecutionResult +from tokenspeed.runtime.grammar.capturable_grammar import setup_grammar_step +from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput +from tokenspeed.runtime.layers.paged_attention import ( + validate_paged_cache_group_ids, +) +from tokenspeed.runtime.sampling.backends.base import SamplingBackend +from tokenspeed.runtime.sampling.dp_sampling_config import ( + DpSamplingRuntimeLimits, + DpSamplingTopology, + resolve_dp_sampling_runtime, + resolve_dp_sampling_support, +) +from tokenspeed.runtime.sampling.sampling_batch_info import SamplingBatchInfo +from tokenspeed.runtime.utils import get_colorful_logger, set_random_seed +from tokenspeed.runtime.utils.common import maybe_inference_mode +from tokenspeed.runtime.utils.env import envs +from tokenspeed.runtime.utils.nvtx import nvtx_range +from tokenspeed.runtime.utils.server_args import ServerArgs + +if TYPE_CHECKING: + from tokenspeed.runtime.layers.attention.backends.base import AttentionBackend + from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool + from tokenspeed.runtime.sampling.sampling_params import SamplingParams + +logger = get_colorful_logger(__name__) + +_DRAFTER_MAPPING = {"EAGLE3": Eagle, "MTP": Eagle, "DFLASH": DFlash} +LOG_MM_TIMING = envs.TOKENSPEED_LOG_MM_TIMING.get() + + +def _eagle_aux_layer_ids(hf_config) -> list[int] | None: + """Draft's eagle_aux_hidden_state_layer_ids (nested or top-level), or None.""" + eagle_config = getattr(hf_config, "eagle_config", None) + if isinstance(eagle_config, dict): + ids = eagle_config.get("eagle_aux_hidden_state_layer_ids") + elif eagle_config is not None: + ids = getattr(eagle_config, "eagle_aux_hidden_state_layer_ids", None) + else: + ids = getattr(hf_config, "eagle_aux_hidden_state_layer_ids", None) + return list(ids) if ids else None + + +def _draft_idle_global_num_tokens_for_step( + step_idx: int, + global_num_tokens: list[int], + global_bs: list[int] | None, +) -> list[int]: + if step_idx == 0 or global_bs is None: + return global_num_tokens + return global_bs + + +PREFILL_GRAPH_DEFAULT_MAX_TOKENS = 2048 + + +def _resolve_prefill_graph_max_tokens(server_args) -> int: + """Largest prefill-graph bucket: explicit value, or min(2048, chunk, kv budget).""" + if server_args.prefill_graph_max_tokens is not None: + return int(server_args.prefill_graph_max_tokens) + cap = PREFILL_GRAPH_DEFAULT_MAX_TOKENS + if server_args.chunked_prefill_size: + cap = min(cap, int(server_args.chunked_prefill_size)) + if server_args.max_total_tokens: + cap = min(cap, int(server_args.max_total_tokens)) + return cap + + +@dataclass +class ModelExecutorConfig: + """ + Scalar configuration for ModelExecutor. + Contains only primitive values — no heavy objects. + Created once via from_server_args() and injected into ModelExecutor. + """ + + # Rank-local graph-padding req-pool index. The C++ scheduler owns real rows + # 1..max_batch_size and row 0 is reserved, so this must sit after the + # scheduler-owned range. + max_req_pool_size: int + output_length: int + enforce_eager: bool + block_size: int + max_num_seqs: int + chunked_prefill_size: int + vocab_size: int + context_len: int + device: str + gpu_id: int + global_rank: int + num_total_pages: int + decode_log_interval: int + cudagraph_capture_sizes: list[int] | None + disable_cuda_graph_padding: bool + max_cudagraph_capture_size: int + model_is_mrope: bool + enable_nan_detection: bool = False + + # ====== DP ========= + data_parallel_size: int = 1 + world_size: int = 1 + world_group: list[int] | None = None + + # ====== SPEC ========= + spec_algo: str | None = None + spec_num_steps: int | None = None + # spec_num_tokens == spec_num_steps + 1 for now (without Tree Attention) + spec_num_tokens: int | None = None + overlap_schedule_depth: int = 0 + dp_sampling: bool = False + dp_sampling_min_bs: int | None = None + use_v4_mtp_paged_metadata: bool = False + + # ====== GRAMMAR ========= + # "none" disables all grammar handling; otherwise the backend name + # (currently only "xgrammar" is implemented). + grammar_backend: str = "xgrammar" + # Force the synchronous eager grammar fallback even on CUDA. For + # parity-testing the captured-grammar path. + disable_capturable_grammar: bool = False + + # ====== PREFILL CUDA GRAPH (breakable) ========= + disable_prefill_graph: bool = False + # Opt-in: > 0 enables the prefill graph and caps the largest token bucket. + prefill_graph_max_tokens: int = 0 + # Explicit bucket list overriding the ladder (see get_prefill_token_buckets). + prefill_graph_capture_sizes: list[int] | None = None + + @staticmethod + def from_server_args( + server_args: ServerArgs, + model_config: ModelConfig, + max_req_pool_size: int, + gpu_id: int, + global_rank: int, + num_total_pages: int, + overlap_schedule_depth: int = 0, + ) -> ModelExecutorConfig: + output_length = ( + server_args.speculative_num_draft_tokens + if server_args.speculative_algorithm + else 1 + ) + rope_parameters = get_rope_parameters(model_config.hf_text_config) + model_is_mrope = bool(rope_parameters and "mrope_section" in rope_parameters) + + return ModelExecutorConfig( + max_req_pool_size=max_req_pool_size, + output_length=output_length, + enforce_eager=server_args.enforce_eager, + block_size=server_args.block_size, + max_num_seqs=server_args.max_num_seqs, + chunked_prefill_size=server_args.chunked_prefill_size, + vocab_size=model_config.vocab_size, + context_len=model_config.context_len, + device=server_args.device, + gpu_id=gpu_id, + global_rank=global_rank, + num_total_pages=num_total_pages, + decode_log_interval=server_args.decode_log_interval, + cudagraph_capture_sizes=server_args.cudagraph_capture_sizes, + disable_cuda_graph_padding=server_args.disable_cuda_graph_padding, + max_cudagraph_capture_size=server_args.max_cudagraph_capture_size, + disable_prefill_graph=bool(server_args.disable_prefill_graph), + prefill_graph_max_tokens=_resolve_prefill_graph_max_tokens(server_args), + prefill_graph_capture_sizes=server_args.prefill_graph_capture_sizes, + model_is_mrope=model_is_mrope, + data_parallel_size=server_args.mapping.attn.dp_size, + world_size=server_args.mapping.world_size, + world_group=server_args.mapping.world_group, + spec_algo=server_args.speculative_algorithm, + spec_num_steps=server_args.speculative_num_steps, + spec_num_tokens=server_args.speculative_num_draft_tokens, + overlap_schedule_depth=overlap_schedule_depth, + dp_sampling=server_args.dp_sampling, + dp_sampling_min_bs=server_args.dp_sampling_min_bs, + enable_nan_detection=server_args.enable_nan_detection, + use_v4_mtp_paged_metadata=model_config.use_v4_mtp_paged_metadata, + grammar_backend=server_args.grammar_backend, + disable_capturable_grammar=server_args.disable_capturable_grammar, + ) + + +class ModelExecutor: + """ + Orchestrates model forward execution. + """ + + def __init__( + self, + config: ModelExecutorConfig, + model_runner: ModelRunner, + attn_backend: AttentionBackend, + token_to_kv_pool: BaseTokenToKVPool, + sampling_backend: SamplingBackend, + draft_model_runner: ModelRunner | None = None, + draft_attn_backend: AttentionBackend | None = None, + draft_token_to_kv_pool: BaseTokenToKVPool | None = None, + mamba_pool: object | None = None, + ): + self.device = config.device + self.config = config + self.model_runner = model_runner + self.sampling_backend = sampling_backend + self.attn_backend = attn_backend + self.token_to_kv_pool = token_to_kv_pool + # Full-attention group mirrored into req_to_page each step (flat+spec). + _group_specs = getattr(token_to_kv_pool, "paged_cache_group_specs", ()) or () + self._flat_full_group_id = next( + ( + str(spec.group_id) + for spec in _group_specs + if getattr(spec, "family", "history") != "state" + and getattr(spec, "retention", None) == "full_history" + ), + None, + ) + self._mirror_idx_cpu: torch.Tensor | None = None + self._mirror_idx_dev: torch.Tensor | None = None + self._mirror_row_buf: torch.Tensor | None = None + self.draft_attn_backend = draft_attn_backend + self.draft_token_to_kv_pool = draft_token_to_kv_pool + self._layerwise_mamba_cow_done = None + + # Must precede CUDA-graph capture: unsupported flat combinations + # (e.g. spec on a backend without flat_spec_capable) would otherwise + # die on a capture-path assert instead of this actionable error. + validate_flat_scheduler_config( + flat_kvcache_ext=scheduler_ext_flat_kvcache(), + paged_cache_groups=_group_specs, + attn_backend=attn_backend, + kv_pool=token_to_kv_pool, + speculative_algorithm=config.spec_algo, + ) + + if config.spec_algo is not None: + # The DFLASH overlap scheduler reserves a fresh draft block per + # decode step a request stays scheduled, including the few steps it + # lingers between finishing and eviction, so peak page count runs + # ~1 page past context_len + spec_num_tokens. Without headroom + # req_to_page overflows and the next draft block's page write goes + # out of bounds, hanging the attention kernel. Pad generously; a few + # int32 columns per request. Non-DFLASH algorithms do not need this. + draft_block_reservation_slack = ( + config.spec_num_tokens * 64 if config.spec_algo == "DFLASH" else 0 + ) + max_num_pages_per_req = ( + config.context_len + + config.spec_num_tokens + + draft_block_reservation_slack + + config.block_size + - 1 + ) // config.block_size + else: + max_num_pages_per_req = ( + config.context_len + config.block_size + ) // config.block_size + + max_bs = config.max_num_seqs // max(config.data_parallel_size, 1) + + self.req_to_page = torch.zeros( + (config.max_req_pool_size + 1, max_num_pages_per_req), + dtype=torch.int32, + device=self.device, + ) + spec_num_tokens = config.spec_num_tokens if config.spec_algo is not None else 1 + self.input_buffers = InputBuffers( + max_bs=max_bs, + max_num_tokens=config.chunked_prefill_size, + page_size=config.block_size, + # token_to_kv_pool allocates size+page_size slots; index `size` is + # the reserved dummy slot (see MHATokenToKVPool._create_buffers). + dummy_kv_slot=0, + state_write_padding_pool_index=config.max_req_pool_size, + device=self.device, + has_mamba=(mamba_pool is not None), + ) + self.runtime_states = RuntimeStates( + req_pool_size=config.max_req_pool_size, + context_len=config.context_len, + vocab_size=config.vocab_size, + device=self.device, + output_length=config.output_length, + mamba_pool=mamba_pool, + ) + # Sized like InputBuffers.max_bs so the padded graph-bucket bs fits. + self.nan_guard = NanGuard.create( + config.enable_nan_detection, + max_bs, + self.device, + ) + if self.config.spec_algo is not None: + DrafterImpl = _DRAFTER_MAPPING[config.spec_algo] + self.drafter = DrafterImpl( + spec_num_tokens=config.spec_num_tokens, + spec_num_steps=config.spec_num_steps, + draft_model_runner=draft_model_runner, + page_size=config.block_size, + runtime_states=self.runtime_states, + input_buffers=self.input_buffers, + req_to_page=self.req_to_page, + attn_backend=draft_attn_backend, + token_to_kv_pool=draft_token_to_kv_pool, + vocab_size=config.vocab_size, + ) + if hasattr(self.drafter, "bind_target_model"): + self.drafter.bind_target_model(self.model_runner.model) + # EAGLE3/MTP share the target's embed + lm_head; DFLASH ships its + # own draft weights, so it must NOT inherit the target's. + if config.spec_algo in ("EAGLE3", "MTP"): + embed, head = self.model_runner.model.get_embed_and_head() + draft_model_runner.model.set_embed_and_head(embed, head) + target_hf = self.model_runner.model_config.hf_config + mm_pad_substitute_id = getattr( + target_hf, "image_token_id", None + ) or getattr(target_hf, "media_placeholder_token_id", None) + if mm_pad_substitute_id is not None and hasattr( + self.drafter, "set_mm_pad_substitute_id" + ): + self.drafter.set_mm_pad_substitute_id(mm_pad_substitute_id) + if config.spec_algo in ("EAGLE3",) and hasattr( + self.model_runner.model, "set_eagle3_layers_to_capture" + ): + # capture the layers the draft was trained on, not the default + aux_layer_ids = _eagle_aux_layer_ids( + draft_model_runner.model_config.hf_config + ) + self.model_runner.model.set_eagle3_layers_to_capture(aux_layer_ids) + if config.spec_algo == "DFLASH": + if not hasattr(self.model_runner.model, "set_dflash_layers_to_capture"): + raise ValueError( + "DFLASH requires the target model to support " + "set_dflash_layers_to_capture." + ) + self.model_runner.model.set_dflash_layers_to_capture( + self.drafter.target_layer_ids + ) + else: + self.drafter = None + + # Single grammar handle: CapturableGrammarExecutor on CUDA (uses + # cudaLaunchHostFunc on a side stream so the xgrammar fill + + # H2D overlap with the forward, and is also CUDA-graph-capturable), + # EagerGrammarBuffers on non-CUDA (synchronous fallback). + # ``disable_capturable_grammar`` forces the eager path on CUDA too + # for parity-testing. + self.grammar_runtime = None + if config.grammar_backend != "none": + from tokenspeed.runtime.grammar.capturable_grammar import ( + CapturableGrammarExecutor, + EagerGrammarBuffers, + ) + + use_captured = ( + current_platform().is_nvidia and not config.disable_capturable_grammar + ) + if use_captured: + self.grammar_runtime = CapturableGrammarExecutor( + max_bs=max_bs, + vocab_size=config.vocab_size, + max_tokens_per_req=spec_num_tokens, + device=self.device, + ) + else: + self.grammar_runtime = EagerGrammarBuffers( + max_bs=max_bs, + vocab_size=config.vocab_size, + max_tokens_per_req=spec_num_tokens, + device=self.device, + ) + + attn_backend.configure_runtime( + sliding_window_size=model_runner.sliding_window_size, + req_to_page=self.req_to_page, + ) + if draft_attn_backend is not None: + draft_attn_backend.configure_runtime( + sliding_window_size=model_runner.sliding_window_size, + req_to_page=self.req_to_page, + ) + + validate_paged_cache_group_ids( + model_runner.model, + token_to_kv_pool.paged_cache_group_specs, + ) + if draft_model_runner is not None and draft_token_to_kv_pool is not None: + validate_paged_cache_group_ids( + draft_model_runner.model, + draft_token_to_kv_pool.paged_cache_group_specs, + ) + + processor = self.model_runner.model.logits_processor + dp_topology = DpSamplingTopology( + tp_rank=processor.tp_rank, + tp_size=processor.tp_size, + tp_group=processor.tp_group, + skip_all_gather=processor.skip_all_gather, + tie_word_embeddings=bool( + getattr(processor.config, "tie_word_embeddings", False) + ), + ) + dp_support = resolve_dp_sampling_support( + requested=self.config.dp_sampling, + drafter_available=self.drafter is not None, + backend_supports_verify=bool( + getattr(self.sampling_backend, "_SUPPORTS_DP_VERIFY", False) + ), + topology=dp_topology, + ) + + lm_head_rows = 0 + if dp_support.enabled: + lm_head_weight = self.model_runner.model.lm_head.weight + if lm_head_weight.ndim < 1: + raise RuntimeError( + "dp_sampling LM head weight must be at least 1D, got " + f"{lm_head_weight.ndim}D" + ) + lm_head_rows = int(lm_head_weight.shape[0]) + dp_runtime_config = resolve_dp_sampling_runtime( + support=dp_support, + lm_head_rows=lm_head_rows, + topology=dp_topology, + limits=DpSamplingRuntimeLimits( + runtime_vocab_size=self.config.vocab_size, + max_num_seqs=config.max_num_seqs, + data_parallel_size=config.data_parallel_size, + num_tokens_per_req=spec_num_tokens, + configured_min_bs=self.config.dp_sampling_min_bs, + device=self.device, + ), + ) + self.dp_sampling_runtime_config = dp_runtime_config + self._last_dp_sampling_route_log: ( + tuple[str, int, bool, int, int, bool, int] | None + ) = None + if dp_runtime_config.enabled: + self.sampling_backend.configure_dp_sampling(dp_runtime_config) + processor.configure_dp_logits_layout(dp_runtime_config) + logger.info( + "Batch-DP spec-verify: requested=%s, infra_supports=%s, enabled=%s " + "min_bs=%s (drafter=%s, backend_supports_dp=%s, " + "tp_size=%s, tp_group=%s)", + dp_support.requested, + dp_support.infra_supports, + dp_support.enabled, + dp_runtime_config.min_bs, + dp_support.drafter_available, + dp_support.backend_supports_verify, + dp_support.tp_size, + dp_support.tp_group_set, + ) + + self._active_multimodal_context = None + self._active_positions_override = None + + self.forward_step = CudaGraphWrapper( + forward_func=self._forward_step, + attn_backend=attn_backend, + token_to_kv_pool=token_to_kv_pool, + input_buffers=self.input_buffers, + config=config, + drafter=self.drafter, + draft_attn_backend=draft_attn_backend, + draft_token_to_kv_pool=draft_token_to_kv_pool, + capturable_grammar=self.capturable_grammar, + eager_grammar_buffers=self.eager_grammar_buffers, + sampling_backend=self.sampling_backend, + runtime_states=self.runtime_states, + ) + + # Breakable prefill (extend) CUDA graphs, the extend-mode analogue of + # the decode wrapper above; captures in __init__, borrowing the decode + # capture stream so all graphs share one mempool-reuse domain. + self.prefill_graph = PrefillGraph( + model_runner=self.model_runner, + attn_backend=attn_backend, + token_to_kv_pool=token_to_kv_pool, + input_buffers=self.input_buffers, + config=config, + req_to_page=self.req_to_page, + drafter=self.drafter, + decode_wrapper=self.forward_step, + ) + + # Encoder CUDA graph: install model-built wrappers by overriding + # modality encoder callables (e.g. ``image_encoder``, ``video_encoder``). + # Multimodal-encoder analogue of ``forward_step``'s ``CudaGraphWrapper``. + self.encoder_graph_wrappers = {} + _mm_model = self.model_runner.model + if ( + hasattr(_mm_model, "make_encoder_cudagraph_wrappers") + and getattr(_mm_model, "is_multimodal_active", True) + and envs.TOKENSPEED_MM_ENABLE_ENCODER_CUDA_GRAPH.get() + and self.model_runner.server_args.mm_attention_backend != "flashinfer_cudnn" + ): + self.encoder_graph_wrappers = _mm_model.make_encoder_cudagraph_wrappers( + _mm_model.mapping + ) + + active_encoder_graph_wrappers = {} + for encoder_attr, wrapper in self.encoder_graph_wrappers.items(): + if not hasattr(_mm_model, encoder_attr): + logger.warning( + "Skipping encoder CUDA graph wrapper for missing attribute %s", + encoder_attr, + ) + continue + setattr(_mm_model, encoder_attr, wrapper) + active_encoder_graph_wrappers[encoder_attr] = wrapper + + self.encoder_graph_wrappers = active_encoder_graph_wrappers + + self.execution_stream = torch.cuda.Stream() + self.log_step = 0 + self._seen_prefill_ids: set[str] = set() + self._prev_decode_bs: int = 0 + self._sentinel_neg1 = torch.tensor(-1, device=self.device, dtype=torch.int64) + if config.model_is_mrope: + mrope_decode_capacity = self.input_buffers.max_num_tokens + # Double-buffered pinned host staging for the decode delta copy. + # Under overlap scheduling the next decode forward is dispatched + # before the previous result is synchronized, so a single reused + # pinned buffer could be refilled by the next step while the prior + # step's ``non_blocking=True`` H2D copy is still reading it (a race + # that corrupts M-RoPE deltas). Ping-pong two buffers so a buffer is + # never overwritten while its copy is in flight (overlap depth 1). + self._mrope_decode_deltas_cpu = [ + self._make_mrope_decode_deltas_cpu(mrope_decode_capacity), + self._make_mrope_decode_deltas_cpu(mrope_decode_capacity), + ] + self._mrope_decode_deltas_cpu_idx = 0 + self._mrope_decode_deltas_buf = torch.zeros( + mrope_decode_capacity, device=self.device, dtype=torch.int64 + ) + else: + self._mrope_decode_deltas_cpu = None + self._mrope_decode_deltas_cpu_idx = 0 + self._mrope_decode_deltas_buf = None + # Decode stats — accumulated from synced results (no GPU sync needed) + self.num_generated_tokens = 0 + self.num_decode_steps = 0 + self.last_decode_stats_tic = time.time() + + set_random_seed(48) + + # Startup has tuned every size class it serves; serving must never autotune. + freeze_autotuning() + + logger.info("ModelExecutor initialized") + + @staticmethod + def _make_mrope_decode_deltas_cpu(size: int) -> torch.Tensor: + try: + return torch.zeros(size, dtype=torch.int64, pin_memory=True) + except RuntimeError: + return torch.zeros(size, dtype=torch.int64) + + @property + def capturable_grammar(self): + """Captured-graph grammar handle, or None on the eager-fallback path. + + Used by ``_forward_step`` to fence the side-stream grammar fill + against the captured forward — those calls only make sense for + the captured flavor of grammar runtime. + """ + from tokenspeed.runtime.grammar.capturable_grammar import ( + CapturableGrammarExecutor, + ) + + return ( + self.grammar_runtime + if isinstance(self.grammar_runtime, CapturableGrammarExecutor) + else None + ) + + @property + def eager_grammar_buffers(self): + """Eager-fallback grammar buffer handle, or None on the captured path.""" + from tokenspeed.runtime.grammar.capturable_grammar import ( + EagerGrammarBuffers, + ) + + return ( + self.grammar_runtime + if isinstance(self.grammar_runtime, EagerGrammarBuffers) + else None + ) + + def _mirror_flat_full_table_into_req_to_page( + self, forward_op, flat_block_tables + ) -> None: + """Flat + spec: scatter the full-attention group's per-batch table + into req_to_page rows, restoring the radix contract for every + legacy consumer (input prep's out_cache_loc kernels, the drafter's + per-step location chains). The flat scheduler never populates + req_to_page itself; column tails zero-fill to the dummy page so + stale longer rows can't leak.""" + if ( + self.drafter is None + or not flat_block_tables + or self._flat_full_group_id is None + ): + return + table = flat_block_tables.get(self._flat_full_group_id) + if table is None: + return + bs = len(forward_op.request_pool_indices) + if self._mirror_idx_cpu is None or self._mirror_idx_cpu.shape[0] < bs: + cap = max(bs, self.input_buffers.max_bs) + self._mirror_idx_cpu = torch.empty(cap, dtype=torch.long, pin_memory=True) + self._mirror_idx_dev = torch.empty( + cap, dtype=torch.long, device=self.device + ) + self._mirror_idx_cpu[:bs] = torch.tensor( + forward_op.request_pool_indices, dtype=torch.long + ) + self._mirror_idx_dev[:bs].copy_(self._mirror_idx_cpu[:bs], non_blocking=True) + idx = self._mirror_idx_dev[:bs] + width = table.shape[1] + max_width = self.req_to_page.shape[1] + if self._mirror_row_buf is None: + self._mirror_row_buf = torch.zeros( + (self.input_buffers.max_bs, max_width), + dtype=self.req_to_page.dtype, + device=self.device, + ) + # Staged rows + index_copy_ (advanced-index setitem costs ~150us dispatch). + rows = self._mirror_row_buf[:bs] + rows[:, :width].copy_(table) + rows[:, :width].clamp_min_( + 0 + ) # -1 column pads -> dummy page 0 (negative locs otherwise) + if width < max_width: + rows[:, width:].zero_() + self.req_to_page.index_copy_(0, idx, rows) + + @nvtx_range("target_forward", color="red") + def _run_target_forward(self, bs: int, ctx: ForwardContext, req_pool_indices): + positions = self._active_positions_override + if positions is None: + if self.config.model_is_mrope: + positions = self.input_buffers.mrope_positions_buf[ + :, : ctx.input_num_tokens + ] + else: + positions = self.input_buffers.positions_buf[: ctx.input_num_tokens] + # Prefill-graph replay when captured for this forward (the decode graph + # replays one level up: it captures the whole _forward_step). The mode + # check is LOAD-BEARING, not an optimization: the decode capture runs + # this dispatch before prefill_graph exists (it is constructed after + # the decode wrapper, whose capture stream it borrows), and decode-mode + # forwards must short-circuit before touching it. + mode = ctx.forward_mode + if ( + mode is not None + and (mode.is_extend() or mode.is_mixed()) + and self.prefill_graph.can_run(ctx, self._active_multimodal_context) + ): + return self.prefill_graph.replay( + ctx, + self.input_buffers.input_ids_buf[: ctx.input_num_tokens], + self._active_multimodal_context, + ) + return self.model_runner.forward( + ctx, + self.input_buffers.input_ids_buf[: ctx.input_num_tokens], + positions, + self.input_buffers.out_cache_loc_buf[: ctx.input_num_tokens], + req_pool_indices=req_pool_indices, + seq_lens=self.input_buffers.seq_lens_buf[:bs], + extend_prefix_lens=self.input_buffers.extend_prefix_lens_buf[ + : ctx.num_extends + ], + multimodal_context=self._active_multimodal_context, + ) + + def _apply_force_single_token_verify( + self, + accept_lengths: torch.Tensor, + row_offset: int, + row_count: int, + decode_input_ids: list[int] | None, + ) -> torch.Tensor: + if decode_input_ids is None or row_count <= 0: + return accept_lengths + force_mask = self.input_buffers.force_single_token_verify_buf[ + row_offset : row_offset + row_count + ] + return torch.where(force_mask, torch.ones_like(accept_lengths), accept_lengths) + + def _cap_accept_to_context_len( + self, + accept_lengths: torch.Tensor, + decode_req_pool_indices: torch.Tensor, + ) -> torch.Tensor: + """Clamp spec-verify accept so committed length never exceeds + ``context_len``. + + ``req_to_page`` is sized for ``context_len + spec_num_tokens`` pages. A + request at the context limit whose ``max_new_tokens`` termination lags a + step can accept past ``context_len``, so its next draft block needs a + page beyond ``req_to_page``'s width — an out-of-bounds access that hangs + the attention kernel. Clamping to the remaining budget keeps the table in + range; the request is still removed a step later. Deterministic in + ``valid_cache_lengths`` / ``accept_lengths``, so no cross-rank divergence. + """ + if accept_lengths.numel() == 0: + return accept_lengths + committed = self.runtime_states.valid_cache_lengths.index_select( + 0, decode_req_pool_indices + ).to(accept_lengths.dtype) + remaining = (self.config.context_len - committed).clamp_(min=0) + # In-place: the drafter reads this same buffer to size its next block. + accept_lengths.copy_(torch.minimum(accept_lengths, remaining)) + return accept_lengths + + def _clamp_committed_to_context_len( + self, + output_lengths: torch.Tensor, + num_extends: int, + bs: int, + ) -> torch.Tensor: + """Return ``output_lengths`` with decode rows clamped so committed KV + length never exceeds ``context_len`` (post-forward, outside the CUDA + graph). + + The clamp must reach BOTH ``_update_runtime_state`` + (``valid_cache_lengths``) and the ``ModelExecutionResult`` that drives + scheduler page reservation, so they stay in lock-step. Hence a FRESH + tensor, not the persistent ``_accept_length_buf``: the verify path also + mirrors accept counts into ``_output_pack_buf``, and an in-place clamp + would leave that mirror (read by the packed-D2H fast path) uncapped, + reserving a draft block past ``req_to_page``'s width and hanging the + kernel. A fresh tensor forces the safe two-D2H fallback. + + Only decode rows ``[num_extends:bs]`` carry an accept delta; prefill rows + pass through. Deterministic, so no cross-rank divergence. + """ + if bs <= num_extends: + return output_lengths + decode_rpi = self.input_buffers.req_pool_indices_buf[num_extends:bs] + committed = self.runtime_states.valid_cache_lengths.index_select( + 0, decode_rpi + ).to(output_lengths.dtype) + remaining = (self.config.context_len - committed).clamp_(min=0) + capped_decode = torch.minimum(output_lengths[num_extends:bs], remaining) + if num_extends == 0: + return capped_decode + return torch.cat([output_lengths[:num_extends], capped_decode]) + + @nvtx_range("sampling", color="yellow") + def _run_sampling( + self, + logits_output: LogitsProcessorOutput, + sampling_info: SamplingBatchInfo, + ctx: ForwardContext, + candidates: torch.Tensor | None = None, + ): + if self.drafter is None: + return self.sampling_backend.sample(logits_output, sampling_info) + + num_extends = ctx.num_extends + num_decodes = ctx.bs - num_extends + + if num_decodes == 0: + return self.sampling_backend.sample(logits_output, sampling_info) + + if num_extends == 0: + output_tokens, accept_lengths = self.sampling_backend.verify( + logits_output, sampling_info, candidates + ) + accept_lengths = self._apply_force_single_token_verify( + accept_lengths, 0, num_decodes, ctx.decode_input_ids + ) + if self.config.spec_algo == "DFLASH": + accept_lengths = self._cap_accept_to_context_len( + accept_lengths, sampling_info.req_pool_indices[:num_decodes] + ) + return output_tokens, accept_lengths + + logits = logits_output.next_token_logits + prefill_out = LogitsProcessorOutput(next_token_logits=logits[:num_extends]) + prefill_tokens, prefill_accept = self.sampling_backend.sample( + prefill_out, sampling_info[:num_extends] + ) + decode_out = LogitsProcessorOutput(next_token_logits=logits[num_extends:]) + decode_tokens, decode_accept = self.sampling_backend.verify( + decode_out, sampling_info[num_extends:], candidates + ) + decode_accept = self._apply_force_single_token_verify( + decode_accept, num_extends, num_decodes, ctx.decode_input_ids + ) + if self.config.spec_algo == "DFLASH": + decode_accept = self._cap_accept_to_context_len( + decode_accept, sampling_info.req_pool_indices[num_extends:] + ) + if ( + prefill_out.next_token_logprobs is not None + and decode_out.next_token_logprobs is not None + ): + logits_output.next_token_logprobs = torch.cat( + [prefill_out.next_token_logprobs, decode_out.next_token_logprobs] + ) + return ( + torch.cat([prefill_tokens, decode_tokens]), + torch.cat([prefill_accept, decode_accept]), + ) + + def _log_dp_sampling_route(self, bs: int, ctx: ForwardContext) -> None: + runtime = self.dp_sampling_runtime_config + if ( + self.config.global_rank != 0 + or not runtime.enabled + or runtime.min_bs is None + or runtime.topology is None + or ctx.forward_mode is None + or not ctx.forward_mode.is_decode() + ): + return + + use_graph = self.forward_step.can_run(bs=bs, ctx=ctx) + effective_bs = self.forward_step.padded_bs(bs=bs, ctx=ctx) if use_graph else bs + tp_size = runtime.topology.tp_size + bucket_bs = ((effective_bs + tp_size - 1) // tp_size) * tp_size + dp_sampling = effective_bs >= runtime.min_bs + route_key = ( + ctx.forward_mode.name, + bs, + use_graph, + effective_bs, + bucket_bs, + dp_sampling, + runtime.min_bs, + ) + if route_key == self._last_dp_sampling_route_log: + return + self._last_dp_sampling_route_log = route_key + logger.debug( + "Batch-DP route: forward_mode=%s bs=%d effective_bs=%d " + "use_graph=%s bucket_bs=%d dp_sampling=%s min_bs=%d", + ctx.forward_mode.name.lower(), + bs, + effective_bs, + use_graph, + bucket_bs, + dp_sampling, + runtime.min_bs, + ) + + @maybe_inference_mode() + def _forward_step( + self, + bs: int, + ctx: ForwardContext, + sampling_info: SamplingBatchInfo, + ): + req_pool_indices = self.input_buffers.req_pool_indices_buf[:bs] + + # Fork grammar onto its side stream so fill + H2D overlap with + # attention/MoE. Rejoined at wait_bitmask() before apply_mask. + if self.capturable_grammar is not None: + n = self.capturable_grammar.max_tokens_per_req + is_spec_verify = n > 1 and ctx.forward_mode.is_decode() + slice_ = ( + self.input_buffers.input_ids_buf[: bs * n] if is_spec_verify else None + ) + self.capturable_grammar.schedule_fill(input_ids_buf_slice=slice_) + + logits_output = self._run_target_forward(bs, ctx, req_pool_indices) + + # Flag NaN per request and sanitize in place, before any sampling kernel. + self.nan_guard.audit_logits(logits_output, ctx) + + candidates = ( + self.drafter.get_candidates(ctx) + if self.config.spec_algo is not None + else None + ) + + if self.capturable_grammar is not None: + self.capturable_grammar.wait_bitmask() + + output_tokens, accept_lengths = self._run_sampling( + logits_output, sampling_info, ctx, candidates + ) + + # Backstop: flag any request whose sampled id falls outside [0, vocab) + # so the output processor can terminate it. Covers sampler/verify kernel + # corruption and DP-sharded steps that audit_logits cannot attribute. + self.nan_guard.merge_oov(output_tokens, ctx, self.runtime_states.vocab_size) + + # Fork sampler-output D2H onto the grammar side stream so the + # next step's build hostfunc can advance the matcher. + if self.capturable_grammar is not None: + self.capturable_grammar.schedule_post_sampler(output_tokens, accept_lengths) + + if self.drafter is not None: + next_round_input_ids = self.drafter.run( + base_ctx=ctx, + logits_output=logits_output, + output_tokens=output_tokens, + accept_lengths=accept_lengths, + ) + # _update_runtime_state skips future_input_map when drafter is + # active — drafter writes the next-round inputs directly. + self.runtime_states.future_input_map[ + self.input_buffers.state_write_req_pool_indices_buf[: ctx.bs] + ] = next_round_input_ids.to(torch.int32) + + output_logprobs = logits_output.next_token_logprobs + return output_tokens, accept_lengths, output_logprobs + + @nvtx_range("update_runtime_state", color="orange") + def _update_runtime_state( + self, + req_pool_indices: torch.Tensor, + output_tokens: torch.Tensor, + accept_lengths: torch.Tensor, + input_lengths: torch.Tensor, + num_extends: int, + ): + """Write output tokens to future_input_map and update cache lengths. + + Must NOT be captured in CUDA graph — these writes are read by the + next iteration's batch prep on the default stream, so they need + explicit stream synchronization (see execute_forward_op). + """ + if self.drafter is None: + # Without drafter, store output tokens for next round. + # With drafter, _forward_step already wrote the drafter's + # next-round input (verified + draft tokens) to future_input_map. + tokens_per_req = self.config.output_length if num_extends == 0 else 1 + next_round_input_ids = output_tokens.to(torch.int32).reshape( + -1, tokens_per_req + ) + self.runtime_states.future_input_map[req_pool_indices, :tokens_per_req] = ( + next_round_input_ids + ) + + bs = req_pool_indices.shape[0] + if num_extends == 0: + deltas = accept_lengths + elif num_extends == bs: + deltas = input_lengths + else: + deltas = torch.cat( + [input_lengths[:num_extends], accept_lengths[num_extends:]] + ) + self.runtime_states.update_valid_cache_length(req_pool_indices, deltas) + + def _build_sampling_info( + self, + bs: int, + sampling_params_list: list[SamplingParams], + ) -> SamplingBatchInfo: + return SamplingBatchInfo( + req_pool_indices=self.input_buffers.req_pool_indices_buf[:bs], + valid_cache_lengths=self.runtime_states.valid_cache_lengths, + is_all_greedy=all(p.top_k <= 1 for p in sampling_params_list), + vocab_size=self.runtime_states.vocab_size, + device=self.device, + ) + + def accumulate_decode_stats(self, results: ModelExecutionResult, bs: int): + """Accumulate decode stats from already-synced results. No GPU sync.""" + self.num_generated_tokens += int(results.output_lengths.sum().item()) + self.num_decode_steps += bs + + @staticmethod + @torch.compile(dynamic=True) + def _compute_mtp_snapshot_indices( + valid_cache_lengths: torch.Tensor, + req_pool_indices: torch.Tensor, + accept_lengths: torch.Tensor, + output_indices: torch.Tensor, + track_indices: torch.Tensor, + sentinel: torch.Tensor, + page_size: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Fused elementwise pipeline computing snapshot src/dst for MTP. + + All operations are batched and fused by torch.compile into a single + Triton kernel (plus the two gathers), eliminating the ~14 individual + elementwise kernel launches of the eager implementation. + """ + new_cl = valid_cache_lengths[req_pool_indices] + old_cl = new_cl - accept_lengths.to(new_cl.dtype) + first_boundary = ((old_cl // page_size) + 1) * page_size + + step_raw = first_boundary - old_cl - 1 + max_col = output_indices.shape[1] - 1 + step = step_raw.clamp(min=0, max=max_col).to(torch.int64) + + bs = req_pool_indices.shape[0] + req_range = torch.arange(bs, device=req_pool_indices.device) + src_raw = output_indices[req_range, step].to(torch.int64) + dst_raw = track_indices.to(torch.int64) + + invalid = ( + (first_boundary > new_cl) + | (dst_raw < 0) + | (src_raw < 0) + | (src_raw == dst_raw) + | (step_raw < 0) + ) + src = torch.where(invalid, sentinel, src_raw) + dst = torch.where(invalid, sentinel, dst_raw) + return src, dst + + def _snapshot_mamba_checkpoints( + self, + accept_lengths: torch.Tensor, + bs: int, + num_extends: int, + ) -> None: + """Snapshot mamba states to checkpoint slots at page boundaries. + + Called after ``_update_runtime_state`` on the execution stream so + ``valid_cache_lengths`` already reflects the accepted tokens. + + Non-MTP (accept_length == 1): + The working slot holds the up-to-date state for the new + cache_length. Pass the kernel page_size so it copies only + when the new length is page-aligned. + + MTP (accept_length > 1): + cache_length may jump over a page boundary. The intermediate + state lives in ``mamba_output_indices[req, step]``. Boundary + detection and source-slot selection are done entirely on GPU + with -1 sentinels so the snapshot kernel skips invalid entries + via its bounds check — no GPU-to-CPU sync, preserving + overlap-schedule pipelining. + """ + if self.runtime_states.mamba_pool is None or num_extends > 0: + return + if not self.input_buffers.has_mamba: + return + + req_pool_indices = self.input_buffers.req_pool_indices_buf[:bs] + track_indices = self.input_buffers.mamba_track_pool_indices_buf[:bs] + page_size = self.config.block_size + dev = req_pool_indices.device + sentinel = self._sentinel_neg1 + + if self.drafter is not None: + # -- MTP path: find the output slot at the crossed boundary -- + backend = getattr( + self.attn_backend, "linear_attn_backend", self.attn_backend + ) + fm = getattr(backend, "forward_metadata", None) + if fm is None: + return + output_indices = fm.mamba_output_indices + if output_indices is None: + return + + src, dst = self._compute_mtp_snapshot_indices( + self.runtime_states.valid_cache_lengths, + req_pool_indices, + accept_lengths[:bs].to(device=dev), + output_indices, + track_indices, + sentinel, + page_size, + ) + + self.runtime_states.snapshot_mamba_checkpoints( + src, + dst, + cache_lengths=None, + page_size=0, + num_valid=bs, + ) + else: + # -- Non-MTP path: working slot IS the up-to-date state -- + src_raw = self.input_buffers.mamba_pool_indices_buf[:bs].to( + device=dev, dtype=torch.int64 + ) + dst_raw = track_indices.to(device=dev, dtype=torch.int64) + + invalid = (src_raw < 0) | (dst_raw < 0) | (src_raw == dst_raw) + src = torch.where(invalid, sentinel, src_raw) + dst = torch.where(invalid, sentinel, dst_raw) + + cache_lengths = self.runtime_states.valid_cache_lengths[req_pool_indices] + self.runtime_states.snapshot_mamba_checkpoints( + src, + dst, + cache_lengths=cache_lengths, + page_size=page_size, + num_valid=bs, + ) + + def flush_mamba_draft_to_working_on_retract(self) -> None: + """Copy accepted draft mamba state -> working slot for all previous-batch requests. + + Called from event_loop when retract WriteBackOps are detected. + Uses the previous decode iteration's input_buffers (still valid since + no new forward has overwritten them). + Runs on execution_stream to respect ordering with previous forward writes. + """ + bs = self._prev_decode_bs + if bs <= 0: + return + + backend = getattr(self.attn_backend, "linear_attn_backend", self.attn_backend) + pool = getattr(backend, "pool", None) + if pool is None: + return + + sentinel = self._sentinel_neg1 + + with torch.cuda.stream(self.execution_stream): + req_pool_indices = self.input_buffers.req_pool_indices_buf[:bs] + working = self.input_buffers.mamba_pool_indices_buf[:bs] + + req = req_pool_indices.to(dtype=torch.int64) + current_input_size = int(pool.current_input_indices.shape[0]) + in_bounds = (req >= 0) & (req < current_input_size) + safe_req = req.clamp(0, current_input_size - 1) + src_raw = pool.current_input_indices[safe_req].to(dtype=torch.int64) + src_raw = torch.where(in_bounds, src_raw, sentinel) + dst_raw = working.to(dtype=torch.int64) + + invalid = ( + (~in_bounds) | (src_raw < 0) | (dst_raw < 0) | (src_raw == dst_raw) + ) + src = torch.where(invalid, sentinel, src_raw) + dst = torch.where(invalid, sentinel, dst_raw) + + self.runtime_states.snapshot_mamba_checkpoints( + src, + dst, + cache_lengths=None, + page_size=0, + num_valid=bs, + ) + + def execute_forward_op_with_log( + self, + forward_op, + sampling_params_list: list[SamplingParams], + num_active_pages: int = 0, + num_cached_pages: int = 0, + num_queue_reqs: int = 0, + dp_global_num_tokens=None, + dp_global_bs=None, + dp_all_decode_or_idle: bool = False, + dp_all_extend: bool = False, + grammar_inputs=None, + multimodal_context=None, + capture_next_input_ids: bool = False, + ) -> ModelExecutionResult: + self.log_step += 1 + + num_extends = forward_op.num_extends() + bs = len(forward_op.request_ids) + is_decode = num_extends <= 0 + + if not is_decode and self.config.global_rank == 0: + mode = "Prefill" if num_extends == bs else "Mix" + total_tokens = sum(forward_op.input_lengths) + cached_tokens = sum( + pl + for rid, pl in zip( + forward_op.request_ids[:num_extends], + forward_op.extend_prefix_lens, + ) + if rid not in self._seen_prefill_ids + ) + if len(self._seen_prefill_ids) > 100_000: + self._seen_prefill_ids.clear() # log-dedup only; bound the growth + self._seen_prefill_ids.update(forward_op.request_ids[:num_extends]) + logger.info( + "%s batch. #new-seq: %s, #new-token: %s, #cached-token: %s, " + "#running-req: %s, #queue-req: %s", + mode, + num_extends, + total_tokens, + cached_tokens, + bs, + num_queue_reqs, + ) + + result = self.execute_forward_op( + forward_op, + sampling_params_list, + dp_global_num_tokens, + dp_global_bs, + dp_all_decode_or_idle, + dp_all_extend, + grammar_inputs=grammar_inputs, + multimodal_context=multimodal_context, + capture_next_input_ids=capture_next_input_ids, + ) + + if is_decode and ( + self.config.global_rank == 0 + and self.log_step % self.config.decode_log_interval == 0 + ): + now = time.time() + gap = now - self.last_decode_stats_tic + gen_throughput = self.num_generated_tokens / gap if gap > 0 else 0 + avg_accept = ( + self.num_generated_tokens / self.num_decode_steps + if self.num_decode_steps > 0 + else 0 + ) + accept_rate = ( + (avg_accept - 1) / self.config.spec_num_steps + if self.config.spec_num_steps + else 0 + ) + num_total_pages = self.config.num_total_pages + page_ratio = ( + num_active_pages / num_total_pages if num_total_pages > 0 else 0 + ) + if self.config.spec_num_steps: + logger.info( + "Decode batch. #running-req: %s, " + "#pages(active/cached/total): %s/%s/%s, " + "page ratio: %.2f, gen throughput (token/s): %.2f, " + "avg_accept_len: %.2f, accept_rate: %.2f, #queue-req: %s", + bs, + num_active_pages, + num_cached_pages, + num_total_pages, + page_ratio, + gen_throughput, + avg_accept, + accept_rate, + num_queue_reqs, + ) + else: + logger.info( + "Decode batch. #running-req: %s, " + "#pages(active/cached/total): %s/%s/%s, " + "page ratio: %.2f, gen throughput (token/s): %.2f, " + "#queue-req: %s", + bs, + num_active_pages, + num_cached_pages, + num_total_pages, + page_ratio, + gen_throughput, + num_queue_reqs, + ) + self.token_to_kv_pool.maybe_log_paged_cache_group_pages() + self.num_generated_tokens = 0 + self.num_decode_steps = 0 + self.last_decode_stats_tic = now + + return result + + def execute_idle_forward( + self, + global_num_tokens: list[int], + global_bs: list[int], + all_decode_or_idle: bool, + ): + """Run a zero-token forward so this rank participates in NCCL collectives. + + Called by the EventLoop when this DP rank has no work but other + ranks do. The MoE all-to-all is a collective that requires ALL + ranks to participate. + """ + graph_forward_mode = ForwardMode.DECODE + ctx = ForwardContext( + attn_backend=self.attn_backend, + token_to_kv_pool=self.token_to_kv_pool, + req_to_page=self.req_to_page, + bs=0, + num_extends=0, + input_num_tokens=0, + forward_mode=graph_forward_mode, + global_num_tokens=global_num_tokens, + global_bs=global_bs, + all_decode_or_idle=all_decode_or_idle, + ) + sampling_info = SamplingBatchInfo( + req_pool_indices=self.input_buffers.req_pool_indices_buf[:0], + valid_cache_lengths=self.runtime_states.valid_cache_lengths, + is_all_greedy=True, + vocab_size=self.runtime_states.vocab_size, + device=self.device, + ) + if self.forward_step.can_run(bs=0, ctx=ctx): + padded_bs = self.forward_step.padded_bs(bs=0, ctx=ctx) + self.input_buffers.fill_dummy_decode_buffers( + batch_size=padded_bs, + total_tokens=padded_bs * self.config.output_length, + ) + # Captured hostfunc pops one entry per replay; push a dummy + # for this idle replay, same as run_once. + if self.capturable_grammar is not None: + self.capturable_grammar.add_batch( + grammars=[None] * padded_bs, bs=padded_bs, has_candidates=False + ) + # IDLE doesn't produce tokens, so no sampler/drafter call here — + # only the model forward, which still participates in collectives. + with nvtx_range("forward_step idle", color="blue"): + self.forward_step( + bs=0, + ctx=ctx, + sampling_info=sampling_info, + req_to_page=self.req_to_page, + ) + return + + # Run model forward with IDLE mode — skips attention but still + # participates in MLP NCCL collectives (dense all-gather, MoE). + ctx.forward_mode = ForwardMode.IDLE + empty = torch.zeros(0, dtype=torch.int32, device=self.device) + self.model_runner.forward( + ctx, + input_ids=empty, + positions=empty, + out_cache_loc=empty, + ) + + # If a drafter is active, its model also has MoE layers that issue + # NCCL collectives. Idle ranks must match those collectives: + # 1 first-step forward + (spec_num_steps - 1) multi-step decode forwards. + if self.drafter is not None: + # DFLASH is a block drafter (idle_forward_steps=1); EAGLE3/MTP + # default to spec_num_steps. Mirror the active rank's per-step + # collective sizing either way. + idle_forward_steps = getattr( + self.drafter, "idle_forward_steps", self.drafter.spec_num_steps + ) + for step_idx in range(idle_forward_steps or 0): + # Mirror active rank's catch-up step: when all non-idle ranks + # are decoding, step 0 sizes collectives from bs/global_bs. + draft_global_num_tokens = _draft_idle_global_num_tokens_for_step( + step_idx, + global_num_tokens, + global_bs, + ) + draft_ctx = ForwardContext( + attn_backend=self.drafter.attn_backend, + token_to_kv_pool=self.drafter.token_to_kv_pool, + req_to_page=self.drafter.req_to_page, + bs=0, + num_extends=0, + input_num_tokens=0, + forward_mode=ForwardMode.IDLE, + global_num_tokens=draft_global_num_tokens, + global_bs=global_bs, + all_decode_or_idle=all_decode_or_idle, + ) + self.drafter.draft_model_runner.forward( + draft_ctx, + input_ids=empty, + positions=empty, + out_cache_loc=empty, + spec_step_idx=step_idx, + ) + + def update_block_table(self, forward_op) -> ModelExecutionResult: + # Update page tables on the default stream before switching to execution stream. + # HostTodevice segment begins + with nvtx_range("update_block_table", color="cyan"): + update_block_table( + forward_op=forward_op, + device=self.device, + req_to_page=self.req_to_page, + ) + + def reset_remote_prefill_mamba_inputs(self, forward_op) -> None: + if self.runtime_states.mamba_pool is None: + return + if not hasattr(self.attn_backend, "reset_current_inputs"): + return + + num_extends = forward_op.num_extends() + if num_extends <= 0: + return + + mamba_indices = list(getattr(forward_op, "mamba_pool_indices", [])) + if not mamba_indices: + return + + req_pool_indices = list(forward_op.request_pool_indices[:num_extends]) + pairs = [ + (int(req_pool_idx), int(mamba_idx)) + for req_pool_idx, mamba_idx in zip( + req_pool_indices, mamba_indices[:num_extends] + ) + if int(mamba_idx) >= 0 + ] + if not pairs: + return + + req_pool_tensor = torch.tensor( + [req_pool_idx for req_pool_idx, _ in pairs], + dtype=torch.int64, + device="cpu", + pin_memory=True, + ).to(self.device, non_blocking=True) + mamba_tensor = torch.tensor( + [mamba_idx for _, mamba_idx in pairs], + dtype=torch.int64, + device="cpu", + pin_memory=True, + ).to(self.device, non_blocking=True) + self.attn_backend.reset_current_inputs(req_pool_tensor, mamba_tensor) + + @staticmethod + def _contains_retracted_decode(forward_op) -> bool: + # FlatForwardOperation stores hist_token_lens for decode rows only; + # non-recovery rows use -1. + return any( + hist_token_len != -1 + for hist_token_len in getattr(forward_op, "hist_token_lens", ()) + ) + + @nvtx_range("reset_valid_cache_length", color="orange") + def reset_valid_cache_length(self, forward_op) -> None: + + num_extends = forward_op.num_extends() + is_prefill = num_extends > 0 + + # Retraction recovery: scheduler pushes -1 per decode op, overriding to + # a real length only on ScheduleDecodeFromRetractedEvent. Decode rows + # may follow prefill rows in a mixed batch, so this cannot be gated on + # pure-decode mode. + has_retract = self._contains_retracted_decode(forward_op) + + # Pure decode without retraction has nothing to do — skip the + # cross-stream wait + stream-context entry entirely. + if not is_prefill and not has_retract: + return + + if has_retract: + hist_token_lens_tensor = torch.tensor( + forward_op.hist_token_lens, + dtype=torch.int32, + device="cpu", + pin_memory=True, + ) + decode_pool_indices = torch.tensor( + forward_op.request_pool_indices[num_extends:], + dtype=torch.int64, + device="cpu", + pin_memory=True, + ) + else: + hist_token_lens_tensor = None + decode_pool_indices = None + + self.execution_stream.wait_stream(torch.cuda.current_stream()) + + with torch.cuda.stream(self.execution_stream): + if is_prefill: + extend_request_pool_indices = torch.tensor( + forward_op.request_pool_indices[:num_extends], + dtype=torch.int64, + device="cpu", + pin_memory=True, + ).to(self.device, non_blocking=True) + + extend_prefix_lens = torch.tensor( + forward_op.extend_prefix_lens, + dtype=torch.int32, + device="cpu", + pin_memory=True, + ).to(self.device, non_blocking=True) + + self.runtime_states.reset_states( + extend_request_pool_indices, extend_prefix_lens + ) + + if hist_token_lens_tensor is not None: + # Apply retraction recovery: override valid_cache_lengths with hist_token_lens + # where the scheduler has specified a non-(-1) value, so that out_cache_loc + # and position IDs are computed against the retracted KV length. + pool_idx_dev = decode_pool_indices.to(self.device, non_blocking=True) + hist_dev = hist_token_lens_tensor.to(self.device, non_blocking=True) + + mask_1d = hist_dev != -1 + vcl = self.runtime_states.valid_cache_lengths[pool_idx_dev] + + self.runtime_states.valid_cache_lengths[pool_idx_dev] = torch.where( + mask_1d, hist_dev, vcl + ) + + def set_layerwise_mamba_cow_done( + self, cow_by_src: dict[int, list[int]] | None + ) -> None: + self._layerwise_mamba_cow_done = ( + { + int(src): {int(dst) for dst in dsts} + for src, dsts in cow_by_src.items() + if dsts + } + if cow_by_src + else None + ) + + def _skip_completed_layerwise_mamba_cow( + self, forward_op, bs: int + ) -> torch.Tensor | None: + cow_done = self._layerwise_mamba_cow_done + self._layerwise_mamba_cow_done = None + if not cow_done or not getattr(self.input_buffers, "has_mamba", False): + return None + cow_src_indices = getattr(forward_op, "mamba_cow_src_indices", None) + working_indices = getattr(forward_op, "mamba_pool_indices", None) + if cow_src_indices is None or working_indices is None: + return None + + cow_src_indices = list(cow_src_indices)[:bs] + working_indices = list(working_indices)[:bs] + skipped_mask = [False] * bs + changed = False + for i, (cow_src, working) in enumerate(zip(cow_src_indices, working_indices)): + cow_src = int(cow_src) + working = int(working) + if working in cow_done.get(cow_src, set()): + cow_src_indices[i] = -1 + skipped_mask[i] = True + changed = True + if not changed: + return None + + self.input_buffers._mamba_cow_src_indices_cpu[:bs].copy_( + torch.as_tensor(cow_src_indices, dtype=torch.int32) + ) + cow_src_buf = self.input_buffers.mamba_cow_src_indices_buf + cow_src_buf[:bs].copy_( + self.input_buffers._mamba_cow_src_indices_cpu[:bs], non_blocking=True + ) + return torch.tensor(skipped_mask, dtype=torch.bool, device=cow_src_buf.device) + + @staticmethod + def _mamba_retract_reset_mask( + mamba_cow_src: torch.Tensor, + bs: int, + skipped_layerwise_cow_mask: torch.Tensor | None, + ) -> torch.Tensor: + reset_mask = mamba_cow_src[:bs] >= 0 + if skipped_layerwise_cow_mask is not None: + reset_mask = reset_mask | skipped_layerwise_cow_mask[:bs].to( + device=reset_mask.device, dtype=torch.bool + ) + return reset_mask + + def _reset_mamba_current_inputs( + self, + *, + num_extends: int, + bs: int, + has_retract: bool, + mamba_pool_indices: torch.Tensor, + mamba_cow_src: torch.Tensor, + skipped_layerwise_cow_mask: torch.Tensor | None, + ) -> None: + if not hasattr(self.attn_backend, "reset_current_inputs"): + return + + if num_extends > 0: + self.attn_backend.reset_current_inputs( + self.input_buffers.req_pool_indices_buf[:num_extends], + mamba_pool_indices[:num_extends], + ) + + if not has_retract: + return + + # FlatForwardOperation places prefill rows first. Restrict recovery to + # the decode suffix so a prefix-cache COW on a prefill row cannot be + # mistaken for a retracted decode. + decode_begin = num_extends + decode_count = bs - decode_begin + if decode_count <= 0: + return + skipped_decode_mask = ( + skipped_layerwise_cow_mask[decode_begin:bs] + if skipped_layerwise_cow_mask is not None + else None + ) + retract_mask = self._mamba_retract_reset_mask( + mamba_cow_src[decode_begin:bs], + decode_count, + skipped_decode_mask, + ) + self.attn_backend.reset_current_inputs( + self.input_buffers.req_pool_indices_buf[decode_begin:bs][retract_mask], + mamba_pool_indices[decode_begin:bs][retract_mask], + ) + + def execute_forward_op( + self, + forward_op, + sampling_params_list: list[SamplingParams], + dp_global_num_tokens=None, + dp_global_bs=None, + dp_all_decode_or_idle: bool = False, + dp_all_extend: bool = False, + grammar_inputs=None, + multimodal_context=None, + capture_next_input_ids: bool = False, + ) -> ModelExecutionResult: + num_extends = forward_op.num_extends() + total_tokens = sum(forward_op.input_lengths) + self._active_multimodal_context = multimodal_context + self._active_positions_override = None + timing_enabled = LOG_MM_TIMING + timing_start = time.perf_counter() if timing_enabled else 0.0 + input_fill_ms = 0.0 + mrope_ms = 0.0 + sampling_prep_ms = 0.0 + forward_step_ms = 0.0 + output_d2h_ms = 0.0 + graph_capable = False + graph_padded_bs = 0 + + with nvtx_range("pre_fill_setup", color="orange"): + has_retract = self._contains_retracted_decode(forward_op) + + # Wait for previous iteration's runtime state updates + # (future_input_map, valid_cache_lengths) on execution_stream to + # complete before reading them. + torch.cuda.current_stream().wait_stream(self.execution_stream) + self.execution_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(self.execution_stream): + bs = len(forward_op.request_ids) + # Outside the graph: in-graph sites only OR into the flag buffer. + self.nan_guard.reset(bs) + # Mirror the full group's flat table into req_to_page (flat+spec). + flat_block_tables = flat_block_tables_from_forward_op( + forward_op, + device=self.device, + num_reqs=bs, + ) + self._mirror_flat_full_table_into_req_to_page(forward_op, flat_block_tables) + decode_input_ids = self.input_buffers.fill_input_buffers( + forward_op=forward_op, + runtime_states=self.runtime_states, + req_to_page=self.req_to_page, + total_tokens=total_tokens, + ) + if timing_enabled: + input_fill_done = time.perf_counter() + input_fill_ms = (input_fill_done - timing_start) * 1000.0 + skipped_layerwise_cow_mask = self._skip_completed_layerwise_mamba_cow( + forward_op, bs + ) + mrope_start = time.perf_counter() if timing_enabled else 0.0 + self._active_positions_override = self._build_mrope_positions_override( + forward_op=forward_op, + multimodal_context=multimodal_context, + total_tokens=total_tokens, + ) + if timing_enabled: + mrope_ms = (time.perf_counter() - mrope_start) * 1000.0 + + forward_mode = ForwardMode.from_num_extends(num_extends, bs) + + if num_extends <= 0: + self._prev_decode_bs = bs + + if self.runtime_states.mamba_pool is not None and ( + num_extends > 0 or has_retract + ): + mamba_pool_indices = self.input_buffers.mamba_pool_indices_buf[:bs] + mamba_cow_src = self.input_buffers.mamba_cow_src_indices_buf[:bs] + self.runtime_states.copy_mamba_states( + mamba_pool_indices, mamba_cow_src, bs + ) + if num_extends > 0: + self.runtime_states.zero_mamba_states( + mamba_pool_indices, + mamba_cow_src, + self.input_buffers.extend_prefix_lens_buf[:num_extends], + num_extends, + ) + self._reset_mamba_current_inputs( + num_extends=num_extends, + bs=bs, + has_retract=has_retract, + mamba_pool_indices=mamba_pool_indices, + mamba_cow_src=mamba_cow_src, + skipped_layerwise_cow_mask=skipped_layerwise_cow_mask, + ) + + grammar_completion = None + + if total_tokens == 0: + # Fully prefix-cached prefill: no tokens to process. + output_tokens = torch.zeros(0, dtype=torch.int32, device=self.device) + output_lengths = torch.zeros(bs, dtype=torch.int32, device=self.device) + output_logprobs = None + else: + gather_ids = None + if num_extends > 0: + num_decodes = bs - num_extends + if self.drafter is not None and num_decodes > 0: + # MIXED + spec: prefill rows pruned to last token, + # decode block kept full at verify width. + num_decode_tokens = num_decodes * self.config.spec_num_tokens + num_prefill_tokens = total_tokens - num_decode_tokens + gather_ids = torch.empty( + num_extends + num_decode_tokens, + dtype=torch.int64, + device=self.device, + ) + gather_ids[:num_extends] = ( + torch.cumsum( + self.input_buffers.input_lengths_buf[:num_extends], + dim=0, + ) + - 1 + ) + gather_ids[num_extends:] = torch.arange( + num_prefill_tokens, + total_tokens, + device=self.device, + dtype=torch.int64, + ) + else: + # EXTEND, MIXED non-spec, or EXTEND + spec: last token + # per request via cumsum. + gather_ids = ( + torch.cumsum( + self.input_buffers.input_lengths_buf[:bs], dim=0 + ) + - 1 + ) + + ctx = ForwardContext( + attn_backend=self.attn_backend, + token_to_kv_pool=self.token_to_kv_pool, + req_to_page=self.req_to_page, + bs=bs, + num_extends=num_extends, + input_num_tokens=total_tokens, + forward_mode=forward_mode, + capture_hidden_mode=( + CaptureHiddenMode.FULL + if self.drafter is not None + else CaptureHiddenMode.NULL + ), + gather_ids=gather_ids, + decode_input_ids=decode_input_ids, + ) + if self.config.data_parallel_size > 1: + if dp_global_num_tokens is None: + raise RuntimeError( + "DP forward metadata must be gathered on CPU by " + "the event loop before model execution." + ) + ctx.global_num_tokens = dp_global_num_tokens + ctx.global_bs = dp_global_bs + ctx.all_decode_or_idle = dp_all_decode_or_idle + ctx.all_extend = dp_all_extend + with nvtx_range("sampling_prep", color="yellow"): + sampling_start = time.perf_counter() if timing_enabled else 0.0 + sampling_info = self._build_sampling_info(bs, sampling_params_list) + grammar_completion = setup_grammar_step( + sampling_info=sampling_info, + bs=bs, + is_spec_decode=self.drafter is not None and num_extends < bs, + spec_num_tokens=self.config.spec_num_tokens or 1, + grammar_inputs=grammar_inputs, + grammar_runtime=self.grammar_runtime, + input_ids_buf=self.input_buffers.input_ids_buf, + grammar_backend=self.config.grammar_backend, + ) + extend_with_prefix = num_extends > 0 and any( + forward_op.extend_prefix_lens + ) + # Flip detection + per-slot scalar scatter + backend-owned + # RNG state refill. Runs OUTSIDE the CUDA graph. Generators + # are now backend-internal (pool-indexed, seeded on flip + # from sp.seed), so the event loop no longer threads them + # through. + self.sampling_backend.prepare_step( + request_ids=forward_op.request_ids, + request_pool_indices=forward_op.request_pool_indices, + sampling_params_list=sampling_params_list, + num_tokens_per_req=self.config.output_length, + ) + if timing_enabled: + sampling_prep_ms = ( + time.perf_counter() - sampling_start + ) * 1000.0 + + with nvtx_range( + f"forward_step ext={num_extends} dec={bs - num_extends}", + color="blue", + ): + mamba_kwargs = ( + { + "mamba_pool_indices": self.input_buffers.mamba_pool_indices_buf[ + :bs + ], + "mamba_cow_src_indices": self.input_buffers.mamba_cow_src_indices_buf[ + :bs + ], + "mamba_branching_seqlens": self.input_buffers.mamba_branching_seqlens_buf[ + :bs + ], + "mamba_track_pool_indices": self.input_buffers.mamba_track_pool_indices_buf[ + :bs + ], + } + if self.input_buffers.has_mamba + else {} + ) + paged_cache_block_tables = paged_cache_block_tables_from_forward_op( + forward_op, + device=self.device, + num_reqs=bs, + ) + # flat_block_tables computed once in pre_fill above. + ( + paged_cache_block_table_base_offsets, + _paged_cache_block_table_base_offset_max, + ) = paged_cache_block_table_base_offsets_from_forward_op( + forward_op, + device=self.device, + num_reqs=bs, + ) + self._log_dp_sampling_route(bs, ctx) + forward_step_start = 0.0 + if timing_enabled: + graph_capable = self.forward_step.can_run(bs, ctx) + graph_padded_bs = ( + self.forward_step.padded_bs(bs, ctx) + if graph_capable + else bs + ) + forward_step_start = time.perf_counter() + output_tokens, output_lengths, output_logprobs = self.forward_step( + bs=bs, + ctx=ctx, + sampling_info=sampling_info, + req_to_page=self.req_to_page, + extend_with_prefix=extend_with_prefix, + extend_prefix_lens=self.input_buffers.extend_prefix_lens_buf[ + :num_extends + ], + extend_prefix_lens_cpu=self.input_buffers.extend_prefix_lens_cpu[ + :num_extends + ], + extend_seq_lens=self.input_buffers.extend_seq_lens_buf[ + :num_extends + ], + extend_seq_lens_cpu=self.input_buffers.extend_seq_lens_cpu[ + :num_extends + ], + paged_cache_block_tables=paged_cache_block_tables, + paged_cache_block_table_base_offsets=( + paged_cache_block_table_base_offsets + ), + flat_block_tables=flat_block_tables, + **mamba_kwargs, + ) + if timing_enabled: + forward_step_ms = ( + time.perf_counter() - forward_step_start + ) * 1000.0 + + if self.config.spec_algo == "DFLASH": + # Clamp the committed-length delta so no request grows past + # context_len. Done here (outside the graph) so it reaches + # both _update_runtime_state and the scheduler page + # reservation; see _clamp_committed_to_context_len. + output_lengths = self._clamp_committed_to_context_len( + output_lengths, num_extends, bs + ) + + # Update runtime state on execution_stream (NOT in the CUDA graph). + self._update_runtime_state( + req_pool_indices=self.input_buffers.req_pool_indices_buf[:bs], + output_tokens=output_tokens, + accept_lengths=output_lengths, + input_lengths=self.input_buffers.input_lengths_buf[:bs], + num_extends=num_extends, + ) + self._snapshot_mamba_checkpoints( + output_lengths, + bs, + num_extends, + ) + + with nvtx_range("output_d2h", color="green"): + output_d2h_start = time.perf_counter() if timing_enabled else 0.0 + next_input_ids = None + if ( + capture_next_input_ids + and self.drafter is not None + and num_extends > 0 + ): + next_input_ids = self.runtime_states.future_input_map.index_select( + 0, self.input_buffers.req_pool_indices_buf[:num_extends] + ).to("cpu", non_blocking=True) + + # Defensive clamp into the valid vocab range (kept from the + # pre-pack path). An out-of-range token id -- e.g. a stale/corrupt + # value surfaced by the intermittent spec-decode decode-state race + # -- would otherwise reach the detokenizer, whose HF + # tokenizer.decode raises a fatal OverflowError on ids outside + # [0, vocab) and tears down the whole server process tree. + # It must run on-GPU *before* the non_blocking D2H: clamping the + # CPU result afterwards would race the in-flight copy. In-place + # (clamp_) so output_tokens keeps aliasing _output_pack_buf and + # the get_packed_output_d2h data_ptr fast-path still fires -- and + # in-place on the forward's inference tensors is only legal inside + # inference mode, so re-enter it (maybe_inference_mode mirrors the + # forward and reduces to no_grad when inference mode is disabled, + # where output_tokens isn't an inference tensor anyway). + vocab_size = self.runtime_states.vocab_size + with maybe_inference_mode(): + output_tokens.clamp_(0, vocab_size - 1) + + packed = self.sampling_backend.get_packed_output_d2h( + output_tokens, output_lengths + ) + if packed is not None: + output_tokens, output_lengths = packed + else: + output_tokens = output_tokens.to("cpu", non_blocking=True) + output_lengths = output_lengths.to("cpu", non_blocking=True) + + if output_logprobs is not None: + output_logprobs = output_logprobs.to("cpu", non_blocking=True) + + output_nan_flags = self.nan_guard.flags_cpu + + copy_event = torch.cuda.Event() + copy_event.record() + if timing_enabled: + output_d2h_ms = (time.perf_counter() - output_d2h_start) * 1000.0 + + if timing_enabled and ( + num_extends > 0 or self.log_step < 64 or self.log_step % 100 == 0 + ): + has_mm = ( + multimodal_context is not None and multimodal_context.has_inputs() + ) + mm_count = 0 + mm_delta_count = 0 + if has_mm: + for mm_input in multimodal_context.mm_inputs: + if mm_input is None: + continue + mm_count += 1 + if mm_input.mrope_position_delta is not None: + mm_delta_count += 1 + logger.info( + "mm_timing forward_execute_ms total=%.3f input_fill=%.3f " + "mrope=%.3f sampling=%.3f forward_step=%.3f output_d2h=%.3f " + "mode=%s bs=%s total_tokens=%s graph=%s padded_bs=%s " + "has_mm=%s mm_count=%s mm_delta_count=%s", + (time.perf_counter() - timing_start) * 1000.0, + input_fill_ms, + mrope_ms, + sampling_prep_ms, + forward_step_ms, + output_d2h_ms, + forward_mode.name, + bs, + total_tokens, + graph_capable, + graph_padded_bs, + has_mm, + mm_count, + mm_delta_count, + ) + + return ModelExecutionResult( + output_tokens=output_tokens, + output_lengths=output_lengths, + output_logprobs=output_logprobs, + copy_event=copy_event, + grammar_completion=grammar_completion, + next_input_ids=next_input_ids, + output_nan_flags=output_nan_flags, + ) + + def write_remote_spec_candidate_ids( + self, req_pool_idx: int, candidate_ids: list[int] + ) -> None: + # Remote spec candidates are CPU materialized; enqueue the H2D copy and + # future_input_map update on execution_stream. The next forward's input + # prep already waits on execution_stream before reading runtime state. + with torch.cuda.stream(self.execution_stream): + self.runtime_states.write_remote_spec_candidate_ids( + req_pool_idx, candidate_ids + ) + + def _expand_mrope_from_input(self, mm_input, seq_len: int) -> torch.Tensor: + # Cache delta expansion for retracted/chunked requests. + if mm_input.mrope_position_delta_repeated_cache is None: + mm_input.mrope_position_delta_repeated_cache = ( + (mm_input.mrope_position_delta - 1).flatten().unsqueeze(0).repeat(3, 1) + ) + return mm_input.mrope_position_delta_repeated_cache + seq_len + + @staticmethod + def _mrope_delta_scalar(mm_input) -> int: + delta = getattr(mm_input, "mrope_position_delta_scalar", None) + if delta is not None: + return int(delta) + tensor = getattr(mm_input, "mrope_position_delta", None) + if tensor is None: + return 0 + delta = int(tensor.flatten()[0].item()) + mm_input.mrope_position_delta_scalar = delta + return delta + + def _build_decode_mrope_positions_override( + self, + forward_op, + mm_inputs, + total_tokens: int, + ) -> torch.Tensor: + if ( + self._mrope_decode_deltas_cpu is None + or self._mrope_decode_deltas_buf is None + ): + raise RuntimeError( + "M-RoPE decode buffers were not initialized for this model" + ) + + base_positions = self.input_buffers.positions_buf[:total_tokens] + # Ping-pong the pinned host staging buffer (see __init__): the previous + # step's non_blocking H2D copy may still be reading the other buffer. + cpu_staging = self._mrope_decode_deltas_cpu[self._mrope_decode_deltas_cpu_idx] + self._mrope_decode_deltas_cpu_idx ^= 1 + token_deltas_cpu = cpu_staging[:total_tokens] + + offset = 0 + has_nonzero_delta = False + for batch_idx, input_len in enumerate(forward_op.input_lengths): + input_len = int(input_len) + if input_len <= 0: + continue + + delta = 0 + mm_input = mm_inputs[batch_idx] if batch_idx < len(mm_inputs) else None + # Honor scalar-only deltas: an upstream payload may set + # mrope_position_delta_scalar while leaving the tensor field + # mrope_position_delta as None (positions precomputed upstream). + # _mrope_delta_scalar handles scalar, tensor, and the absent case + # (returns 0), so call it whenever an mm_input is present. + if mm_input is not None: + delta = self._mrope_delta_scalar(mm_input) + has_nonzero_delta = has_nonzero_delta or delta != 0 + + token_deltas_cpu[offset : offset + input_len].fill_(delta) + offset += input_len + + if offset != total_tokens: + token_deltas_cpu[offset:total_tokens].zero_() + + if has_nonzero_delta: + token_deltas = self._mrope_decode_deltas_buf[:total_tokens] + token_deltas.copy_(token_deltas_cpu, non_blocking=True) + mrope_base = base_positions + token_deltas + else: + mrope_base = base_positions + + self.input_buffers.mrope_positions_buf[:, :total_tokens].copy_( + mrope_base.unsqueeze(0).expand(3, -1) + ) + return self.input_buffers.mrope_positions_buf[:, :total_tokens] + + def _build_mrope_positions_override( + self, + forward_op, + multimodal_context, + total_tokens: int, + ) -> torch.Tensor | None: + if not self.config.model_is_mrope or total_tokens == 0: + return None + + is_prefill = forward_op.num_extends() > 0 + base_positions = self.input_buffers.positions_buf[:total_tokens] + mm_inputs = ( + multimodal_context.mm_inputs + if multimodal_context is not None and multimodal_context.has_inputs() + else [] + ) + if not is_prefill: + return self._build_decode_mrope_positions_override( + forward_op=forward_op, + mm_inputs=mm_inputs, + total_tokens=total_tokens, + ) + + pos_chunks = torch.split(base_positions, list(forward_op.input_lengths), dim=0) + mrope_chunks = [] + for batch_idx, base_chunk in enumerate(pos_chunks): + mm_input = mm_inputs[batch_idx] if batch_idx < len(mm_inputs) else None + # Fall back to linear only when there is neither a per-token mrope table + # nor a transferred scalar delta. A decode-only mm_input may carry just + # the delta (post-image decode positions = base+delta); it must skip the + # fallback and take the base+delta branch below. + if mm_input is None or ( + mm_input.mrope_positions is None + and mm_input.mrope_position_delta is None + ): + mrope_chunks.append(base_chunk.unsqueeze(0).expand(3, -1)) + continue + + if ( + is_prefill + and mm_input.mrope_positions is not None + and batch_idx < len(forward_op.extend_prefix_lens) + ): + start = int(forward_op.extend_prefix_lens[batch_idx]) + end = start + int(forward_op.input_lengths[batch_idx]) + positions = mm_input.mrope_positions[:, start:end] + if positions.numel() != 0: + mrope_chunks.append( + positions.to(device=self.device, dtype=torch.int64) + ) + continue + if base_chunk.numel() == 1: + seq_len = int(base_chunk[-1].item()) + 1 + mrope_chunks.append( + self._expand_mrope_from_input(mm_input, seq_len).to( + device=self.device, dtype=torch.int64 + ) + ) + continue + + delta = mm_input.mrope_position_delta + if delta is None: + delta = torch.zeros(1, dtype=torch.int64) + delta = delta.flatten()[0].to(device=self.device, dtype=torch.int64) + # Decode positions need (mrope_delta - 1) + seq_len. positions_buf + # already stores the per-token zero-based position (seq_len - 1 for + # decode), so this is the same value without a GPU-to-CPU sync. + mrope_chunks.append((base_chunk + delta).unsqueeze(0).expand(3, -1)) + + mrope_positions = torch.cat(mrope_chunks, dim=1).contiguous() + self.input_buffers.mrope_positions_buf[:, :total_tokens].copy_(mrope_positions) + return self.input_buffers.mrope_positions_buf[:, :total_tokens] diff --git a/python/tokenspeed/runtime/execution/model_runner.py b/python/tokenspeed/runtime/execution/model_runner.py new file mode 100644 index 0000000..5735c3f --- /dev/null +++ b/python/tokenspeed/runtime/execution/model_runner.py @@ -0,0 +1,167 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import inspect +from typing import TYPE_CHECKING + +import torch + +from tokenspeed.runtime.execution.weight_loader import WeightLoader +from tokenspeed.runtime.layers.moe.utils import initialize_moe_config +from tokenspeed.runtime.utils import get_colorful_logger +from tokenspeed.runtime.utils.env import global_server_args_dict_update +from tokenspeed.runtime.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter + +if TYPE_CHECKING: + from tokenspeed.runtime.configs.model_config import ModelConfig + from tokenspeed.runtime.execution.context import ForwardContext + from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput + from tokenspeed.runtime.multimodal.inputs import MultimodalForwardContext + from tokenspeed.runtime.utils.server_args import ServerArgs + +logger = get_colorful_logger(__name__) + + +class ModelRunner: + def __init__( + self, + # Configuration + model_config: ModelConfig, + server_args: ServerArgs, + gpu_id: int, + global_rank: int, + is_draft_worker: bool = False, + ): + """Initialize ModelRunner with injected dependencies.""" + # Store configuration + self.model_config = model_config + self.server_args = server_args + self.device = server_args.device + self.gpu_id = gpu_id + self.global_rank = global_rank + self.mapping = server_args.mapping + self.is_generation = model_config.is_generation + self.is_multimodal = model_config.is_multimodal + self.is_draft_worker = is_draft_worker + self.mambaish_config = getattr(model_config, "mambaish_config", None) + self.is_hybrid_gdn = getattr(model_config, "is_hybrid_gdn", False) + self.sliding_window_size = getattr( + model_config.hf_config, "sliding_window", None + ) + + draft_moe_override = ( + self.is_draft_worker + and server_args.draft_moe_backend is not None + and server_args.draft_moe_backend != server_args.moe_backend + ) + if draft_moe_override: + saved_moe_backend = server_args.moe_backend + server_args.moe_backend = server_args.draft_moe_backend + + # Auto-detect FP8 KV cache from checkpoint quant config (e.g. NVFP4 models + # with kv_cache_quant_algo: "FP8" in hf_quant_config.json). + if server_args.kv_cache_dtype == "auto": + quant_cfg = model_config._parse_quant_hf_config() + if quant_cfg is not None: + kv_algo = quant_cfg.get("kv_cache_quant_algo") + if isinstance(kv_algo, str) and kv_algo.upper() == "FP8": + server_args.kv_cache_dtype = "fp8_e4m3" + logger.info( + "Auto-detected kv_cache_dtype=fp8_e4m3 from checkpoint " + "quant config (kv_cache_quant_algo=%s)", + kv_algo, + ) + + global_server_args_dict_update(server_args) + initialize_moe_config(server_args) + + self.memory_saver_adapter = TorchMemorySaverAdapter.create( + enable=server_args.enable_memory_saver + ) + self.load_model() + if draft_moe_override: + server_args.moe_backend = saved_moe_backend + global_server_args_dict_update(server_args) + initialize_moe_config(server_args) + + def load_model(self): + self.model = WeightLoader.load_model( + model_config=self.model_config, + server_args=self.server_args, + device=self.device, + gpu_id=self.gpu_id, + memory_saver_adapter=self.memory_saver_adapter, + ) + self._model_forward_accepts_spec_step_idx = self._forward_accepts_kwarg( + self.model, "spec_step_idx" + ) + + @staticmethod + def _forward_accepts_kwarg(model, name: str) -> bool: + try: + parameters = inspect.signature(model.forward).parameters + except (TypeError, ValueError): + return False + + return name in parameters + + def forward( + self, + ctx: ForwardContext, + input_ids: torch.Tensor, + positions: torch.Tensor, + out_cache_loc: torch.Tensor, + req_pool_indices: torch.Tensor | None = None, + seq_lens: torch.Tensor | None = None, + extend_prefix_lens: torch.Tensor | None = None, + captured_hidden_states: torch.Tensor | None = None, + input_embeds: torch.Tensor | None = None, + multimodal_context: MultimodalForwardContext | None = None, + spec_step_idx: int | None = None, + ) -> LogitsProcessorOutput: + kwargs = {} + if req_pool_indices is not None: + kwargs["req_pool_indices"] = req_pool_indices + if seq_lens is not None: + kwargs["seq_lens"] = seq_lens + if extend_prefix_lens is not None: + kwargs["extend_prefix_lens"] = extend_prefix_lens + if not self.is_generation: + kwargs["get_embedding"] = True + if captured_hidden_states is not None: + kwargs["captured_hidden_states"] = captured_hidden_states + if input_embeds is not None: + kwargs["input_embeds"] = input_embeds + if multimodal_context is not None: + kwargs["multimodal_context"] = multimodal_context + if spec_step_idx is not None and getattr( + self, "_model_forward_accepts_spec_step_idx", False + ): + kwargs["spec_step_idx"] = spec_step_idx + + return self.model.forward( + ctx, + input_ids, + positions, + out_cache_loc, + **kwargs, + ) diff --git a/python/tokenspeed/runtime/execution/nan_guard.py b/python/tokenspeed/runtime/execution/nan_guard.py new file mode 100644 index 0000000..7f7a32f --- /dev/null +++ b/python/tokenspeed/runtime/execution/nan_guard.py @@ -0,0 +1,155 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Per-request NaN containment for the model executor. + +Models and kernels cannot be assumed 100% NaN-free. This guard records +*which* requests produced NaN logits (or an out-of-vocab token id), +sanitizes the logits in place, and ships a per-request flag tensor to the +CPU where the output processor terminates the flagged requests with +``ABORT_CODE.NumericalError``. + +Design constraints (all hold by construction): + +- **Graph-safe / no sync.** Fixed-shape device ops on a persistent flag + buffer; flags OR in-graph and are zeroed once per step outside it, so + multi-cycle decode graphs accumulate correctly. +- **Near-zero cost.** Detection is one fused ``amax`` reduction over the + logits (NaN propagates through ``amax``; no ``[rows, vocab]`` mask) plus + ops on ``[bs]``-sized vectors; sanitize is one ``nan_to_num_``. +- **Rank-consistent.** OOV flags derive from already-broadcast token ids; + logits flags rely on the bit-identical-logits-per-rank assumption the + conditional sampling broadcast already depends on. +- **Zero branching at call sites.** ``NanGuard.create`` returns a no-op + singleton when disabled. + +Limitation: with Batch-DP spec-verify sampling the logits arrive sharded +(``logits_layout_plan is not None``), so logits attribution is skipped for +those steps; sanitize still applies and the OOV backstop still covers the +gathered full-batch ids. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from tokenspeed.runtime.execution.forward_context import ForwardContext + from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput + +# Replacement for NaN / +-inf logits (fp32). +-1e30 leaves headroom so a +# later temperature division cannot overflow back to inf; matches SGLang. +_NEG_SANITIZED = -1e30 +_POS_SANITIZED = 1e30 + + +class NanGuard: + """Tracks per-request numerical corruption across one executor step. + + Lifecycle per ``execute_forward_op``:: + + guard.reset(bs) # outside the graph + ... per forward cycle (in-graph): + guard.audit_logits(logits_output, ctx) # pre-sampling + guard.merge_oov(tokens, ctx, vocab_size) # OOV backstop + flags = guard.flags_cpu # with the output D2H batch + """ + + def __init__(self, max_bs: int, device: torch.device | str) -> None: + self.flags = torch.zeros((max_bs,), dtype=torch.int32, device=device) + self._bs = 0 + + @classmethod + def create(cls, enabled: bool, max_bs: int, device) -> NanGuard: + return cls(max_bs, device) if enabled else _DISABLED + + def reset(self, bs: int) -> None: + """Zero the flags and pin this step's batch size; call outside the graph.""" + self._bs = bs + self.flags.zero_() + + def audit_logits( + self, logits_output: LogitsProcessorOutput, ctx: ForwardContext + ) -> None: + """Flag requests with NaN logits, then sanitize the logits in place. + + Must run before sampling and before grammar vocab masks / logit_bias, + so their legitimate ``-inf`` entries survive sanitize. + """ + logits = logits_output.next_token_logits + if logits_output.logits_layout_plan is None: + self._or_per_request(torch.isnan(logits.amax(dim=-1)), ctx) + torch.nan_to_num_( + logits, nan=_NEG_SANITIZED, posinf=_POS_SANITIZED, neginf=_NEG_SANITIZED + ) + + def merge_oov( + self, output_tokens: torch.Tensor, ctx: ForwardContext, vocab_size: int + ) -> None: + """Backstop: flag requests whose sampled ids fall outside [0, vocab). + + Catches corruption past the logits (sampler/verify kernel output) and + covers DP-sharded steps. Token ids are already rank-synced here. + """ + self._or_per_request((output_tokens < 0) | (output_tokens >= vocab_size), ctx) + + @property + def flags_cpu(self) -> torch.Tensor | None: + """Async D2H of this step's flags (order with the copy event).""" + return self.flags[: self._bs].to("cpu", non_blocking=True) + + def _or_per_request(self, rows: torch.Tensor, ctx: ForwardContext) -> None: + """OR a per-row bool vector into per-request flags. + + Row layout mirrors ``_run_sampling``: ``[num_extends]`` extend rows, + then ``num_decodes * n`` decode/verify rows. + """ + ne = ctx.num_extends + nd = ctx.bs - ne + if ne > 0: + self.flags[:ne] |= rows[:ne].to(torch.int32) + if nd > 0: + n = (rows.shape[0] - ne) // nd + self.flags[ne : ctx.bs] |= rows[ne:].view(nd, n).any(dim=-1).to(torch.int32) + + +class _DisabledNanGuard(NanGuard): + """No-op stand-in so call sites need no enabled-checks.""" + + def __init__(self) -> None: # no buffer + pass + + def reset(self, bs: int) -> None: + pass + + def audit_logits(self, logits_output, ctx) -> None: + pass + + def merge_oov(self, output_tokens, ctx, vocab_size) -> None: + pass + + @property + def flags_cpu(self) -> None: + return None + + +_DISABLED = _DisabledNanGuard() diff --git a/python/tokenspeed/runtime/execution/prefill_graph.py b/python/tokenspeed/runtime/execution/prefill_graph.py new file mode 100644 index 0000000..796e669 --- /dev/null +++ b/python/tokenspeed/runtime/execution/prefill_graph.py @@ -0,0 +1,656 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Breakable CUDA graphs for prefill (extend) forwards. + +:class:`PrefillGraph` holds one breakable graph per padded token bucket +(captured from a dummy bs=1 extend batch). The embedding lookup stays OUTSIDE +the captured region: graphs start from a static input-embeds buffer, filled at +replay by an eager ``embed_tokens`` gather (text) or by precomputed merged +embeddings (multimodal, via the model's ``multimodal_input_embeds`` seam). +Constructed after the decode +:class:`~tokenspeed.runtime.execution.cuda_graph_wrapper.CudaGraphWrapper`, +borrowing its capture stream; buckets share one private mempool, deliberately +not the decode graphs' pool (see :meth:`capture`). At serving time +the executor's target-forward dispatch is a flat +three-way -- decode & captured replays the decode graph (one level up, since +it captures the whole step), prefill & captured replays here (:meth:`can_run` +/ :meth:`replay`), everything else runs the eager model forward. + +Unlike decode (whole forward captured, keyed by batch size), the captured +region here is purely token-shaped compute keyed by total token count: +attention runs as an eager break (see +:mod:`tokenspeed.runtime.execution.breakable_cuda_graph`), so one graph per +bucket serves any batch size at that token count, and a replayed forward is +finished with the model's eager logits tail. +""" + +from __future__ import annotations + +import bisect +from contextlib import contextmanager +from typing import TYPE_CHECKING, NamedTuple + +import torch + +from tokenspeed.runtime.execution.breakable_cuda_graph import ( + BreakableCapture, + active_forward, +) +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.execution.forward_batch_info import ( + CaptureHiddenMode, + ForwardMode, +) +from tokenspeed.runtime.layers.logits_processor import LogitsMetadata +from tokenspeed.runtime.utils import get_colorful_logger +from tokenspeed.runtime.utils.common import maybe_inference_mode + +logger = get_colorful_logger(__name__) + +if TYPE_CHECKING: + from tokenspeed.runtime.execution.cuda_graph_wrapper import CudaGraphWrapper + from tokenspeed.runtime.execution.input_buffer import InputBuffers + from tokenspeed.runtime.execution.model_executor import ModelExecutorConfig + from tokenspeed.runtime.layers.attention.backends.base import AttentionBackend + + +# Smallest prefill bucket; below this, denser rungs would only add capture time. +PREFILL_BUCKET_FLOOR: int = 16 + +# Relative rung spacing (largest pow2 <= size/8), bounding the padded tail at ~12.5%. +PREFILL_BUCKET_STEP_DIVISOR: int = 8 + +# Absolute rung-spacing cap, bounding the worst case at the top of the ladder. +PREFILL_BUCKET_MAX_STEP: int = 512 + + +def get_prefill_token_buckets(config: ModelExecutorConfig) -> list[int]: + """Padded token-count buckets to capture for the breakable prefill graph. + + Unlike decode (keyed by batch size), the breakable prefill graph captures + pure token-shaped compute, so it is keyed by total token count. A live extend + forward is padded up to the smallest bucket >= its token count; forwards above + the largest bucket run eager. + + Returns an empty list (graph disabled) when ``disable_prefill_graph`` is set or + ``prefill_graph_max_tokens <= 0``. The largest bucket is clamped to the + chunked-prefill size: the scheduler's per-forward token budget + (``max_scheduled_tokens`` = chunked-prefill size) covers extends AND any fused + decode rows -- with mixed batching, decodes are scheduled first and each + decrements the budget, and the prefill chunk is sized to what remains + (scheduler ``newForwardOperation``/``push_op``) -- so no forward, mixed or + pure, ever exceeds the chunk. No headroom above it is needed. + + The default ladder bounds RELATIVE padding waste: a forward pads its graphed + compute to the next bucket, so what matters is the gap as a fraction of the + size -- a flat stride is needlessly coarse for short prompts and needlessly + dense at the top. Each bucket's step is the largest power of two <= size/8 + (padded tail at most ~12.5% anywhere on the ladder), floored at 16 tokens and + capped at 512 so the absolute worst case stays bounded at the top end. Dense + ladders are cheap: all captures share one stream + mempool, so graph memory + is ~the largest bucket's peak regardless of bucket count (see + ``BreakableCapture``); the remaining cost is ~0.5s of startup capture per + bucket. + + ``prefill_graph_capture_sizes`` overrides the ladder with an explicit list + (mirroring decode's ``cudagraph_capture_sizes``) -- e.g. a short list for + faster startup on dev boots; sizes are clamped to the largest bucket. + + Args: + config: The model-executor config carrying ``disable_prefill_graph``, + ``prefill_graph_max_tokens``, ``prefill_graph_capture_sizes`` and + ``chunked_prefill_size``. + + Returns: + Sorted ascending list of token-bucket sizes (possibly empty). + """ + max_tokens = int(config.prefill_graph_max_tokens or 0) + if config.disable_prefill_graph or max_tokens <= 0: + return [] + chunk = int(config.chunked_prefill_size or 0) + if chunk > 0: + max_tokens = min(max_tokens, chunk) + explicit = config.prefill_graph_capture_sizes + if explicit: + buckets = {int(b) for b in explicit if 0 < int(b) <= max_tokens} + buckets.add(max_tokens) + return sorted(buckets) + buckets = [] + size = min(PREFILL_BUCKET_FLOOR, max_tokens) + while size < max_tokens: + buckets.append(size) + size += _prefill_bucket_step(size) + buckets.append(max_tokens) + return sorted(set(buckets)) + + +def _prefill_bucket_step(size: int) -> int: + """Distance from bucket ``size`` to the next rung. + + The largest power of two <= ``size / PREFILL_BUCKET_STEP_DIVISOR`` (so the + padded tail stays within ~1/8 of the real token count), clamped between + ``PREFILL_BUCKET_FLOOR`` and ``PREFILL_BUCKET_MAX_STEP``. + """ + relative = size // PREFILL_BUCKET_STEP_DIVISOR + if relative <= PREFILL_BUCKET_FLOOR: + return PREFILL_BUCKET_FLOOR + largest_pow2 = 1 << (relative.bit_length() - 1) + return min(largest_pow2, PREFILL_BUCKET_MAX_STEP) + + +class CapturedForward(NamedTuple): + """A bucket's captured inner-forward outputs (stable pool addresses).""" + + # Final hidden states with shape [bucket, hidden]; padded tail is garbage. + hidden_states: torch.Tensor + + # Aux hidden states for drafting, each [bucket, hidden]; None when mode is NULL. + aux_hidden_states: list[torch.Tensor] | None + + def sliced(self, num_tokens: int) -> tuple[torch.Tensor, list[torch.Tensor] | None]: + """The leading real-token rows, in the (hidden, aux) shape callers expect.""" + hidden = self.hidden_states[:num_tokens] + if self.aux_hidden_states is None: + return hidden, None + return hidden, [a[:num_tokens] for a in self.aux_hidden_states] + + +class PrefillGraph: + """The breakable prefill (extend) CUDA graphs. + + A pure graph object -- :meth:`can_run` / :meth:`replay` -- holding no + reference to any other component. Constructed AFTER the decode + ``CudaGraphWrapper`` and captures in ``__init__`` like it: the decode + wrapper is used transiently for its capture stream and dummy paged-cache + tables, not kept. (The executor's target-forward dispatch therefore mode- + checks before touching this object -- decode capture runs that dispatch + while this object does not exist yet.) The dispatch checks :meth:`can_run` + and calls :meth:`replay`; the eager path stays a direct + ``model_runner.forward`` call at that call site. Capture failure degrades + to eager -- world-agreed, so DP/TP ranks stay in lockstep. + + Args: + model_runner: The target ModelRunner. Supplies the loaded model + (multimodal wrappers are unwrapped internally: the graph wraps the + nested ``language_model``'s text transformer, image prefills run + eager) and ``is_generation`` (embedding models run eager). + attn_backend: Backend whose extend metadata the dummy capture batch sets. + token_to_kv_pool: KV pool the dummy batch points at (reserved dummy slot). + input_buffers: The shared static input buffers the graphs read from. + config: Model-executor config (buckets, DP/world topology, device). + req_to_page: Request page table; row 0 backs the dummy capture request. + drafter: If present, aux-hidden capture (EAGLE3/MTP) is baked into the + captured graphs. + """ + + def __init__( + self, + model_runner, + attn_backend: AttentionBackend, + token_to_kv_pool, + input_buffers: InputBuffers, + config: ModelExecutorConfig, + req_to_page: torch.Tensor | None, + drafter=None, + decode_wrapper: CudaGraphWrapper | None = None, + num_warmup: int = 3, + ) -> None: + model = model_runner.model if model_runner is not None else None + # Multimodal seam: models whose multimodal path is embeds-only expose + # multimodal_input_embeds; others (e.g. deepstack) replay text only. + self._multimodal_input_embeds = getattr(model, "multimodal_input_embeds", None) + self.text_model = ( + model.language_model if hasattr(model, "language_model") else model + ) + self.inner_model = getattr(self.text_model, "model", None) + # Embedding runs eagerly OUTSIDE the graphs (see capture); the graphs + # read a static input-embeds buffer instead of gathering from input_ids. + self._embed_tokens = getattr(self.inner_model, "embed_tokens", None) + self._input_embeds_buf: torch.Tensor | None = None + self.attn_backend = attn_backend + self.token_to_kv_pool = token_to_kv_pool + self.input_buffers = input_buffers + self.config = config + self.req_to_page = req_to_page + self.drafter = drafter + self.num_warmup = num_warmup + self.dp_size = config.data_parallel_size + + self.capture_buckets = get_prefill_token_buckets(config) + self.disable = ( + config.enforce_eager + or config.disable_prefill_graph + or not self.capture_buckets + or self.inner_model is None + or self._embed_tokens is None + or model_runner is None + or not model_runner.is_generation + # DP replay decisions must come from replicated state, and a + # forward's multimodal-ness is rank-local: one rank running its mm + # prefill eager while text-only peers replay desyncs the EP + # collectives. Until the DP metadata gather carries a multimodal + # flag, keep the graph off for multimodal models under DP. + or (config.data_parallel_size > 1 and model_runner.is_multimodal) + ) + + self._ctx: ForwardContext | None = None + self._pool = None + self._engaged_logged: set[str] = set() + # Aux-capture mode baked into the graphs; mismatched live forwards run eager. + self._captured_hidden_mode = None + # One captured graph + bucket-sized output per padded token bucket. + self._captures: dict[int, BreakableCapture] = {} + self._outputs: dict[int, CapturedForward] = {} + + if not self.disable: + self.capture(decode_wrapper) + + # ------------------------------------------------------------------ + # Graph capture + # ------------------------------------------------------------------ + + def capture(self, decode_wrapper: CudaGraphWrapper | None = None) -> None: + """Capture one breakable graph per token bucket (no-op when disabled). + + Called from ``__init__``; ``decode_wrapper`` supplies the shared + capture stream and dummy paged-cache block tables (used here only, + not stored). Buckets share one PRIVATE mempool (first capture + allocates it), so graph memory stays ~the largest bucket's peak -- + but never the decode graphs' pool: eager ops cache raw pointers to + buffers they lazily allocated inside a decode capture (flashinfer's + trtllm-gen MoE runner), and a prefill capture reusing those freed + blocks means every replay rewrites them, corrupting the next eager + call (IMA; A/B-proven on qwen3.5 MTP). + + Runs under inference mode like serving forwards (in-place updates on + inference-mode model state buffers are only legal there). OOM fails + the boot LOUDLY (the graph pool did not fit next to weights + KV + cache; the operator decides: free headroom, lower + ``--prefill-graph-max-tokens``, or 0 to disable). Any other failure + means the dummy-batch machinery doesn't cover this model family yet: + degrade to eager prefill instead of crashing the server, and agree on + that across the world (a MIN all-reduce over the success flag) -- + replay force-sets ``global_num_tokens`` on every rank, so one eager + rank among replaying peers diverges the token counts and deadlocks + the next collective. + """ + if self.disable: + return + weight = self._embed_tokens.weight + self._input_embeds_buf = torch.zeros( + max(self.capture_buckets), + weight.shape[1], + dtype=weight.dtype, + device=weight.device, + ) + captured_ok = True + try: + with maybe_inference_mode(): + self._capture_all_buckets(decode_wrapper) + except torch.cuda.OutOfMemoryError: + logger.error( + "Prefill graph capture ran out of GPU memory. Free up " + "--gpu-memory-utilization headroom, lower " + "--prefill-graph-max-tokens (default %d), or set it to 0 to " + "disable the prefill graph.", + 2048, + ) + raise + except (NotImplementedError, AttributeError, KeyError, RuntimeError) as exc: + logger.warning( + "Prefill graph capture failed (%s: %s); falling back to eager " + "prefill. This model family may need dedicated dummy-batch support.", + type(exc).__name__, + exc, + ) + captured_ok = False + if not self._capture_unanimous(captured_ok): + self.disable = True + + def _capture_all_buckets(self, decode_wrapper: CudaGraphWrapper | None) -> None: + for bucket in sorted(self.capture_buckets, reverse=True): + self._ctx = self._make_dummy_batch(bucket, decode_wrapper) + self._land_input_embeds( + self._embed_tokens(self.input_buffers.input_ids_buf[:bucket]), bucket + ) + self._captured_hidden_mode = self._ctx.capture_hidden_mode + # Breaks record the ambient dummy ctx; it is rebound live at replay. + try: + with active_forward(self._ctx): + self._capture_bucket(bucket, decode_wrapper) + finally: + self._ctx = None + if self.config.global_rank == 0: + sample = next(iter(self._captures.values()), None) + logger.info( + "prefill breakable graph: captured buckets %s (segments=%d, eager " + "attention breaks)", + sorted(self._captures), + sample.num_segments if sample is not None else 0, + ) + + def _capture_bucket( + self, bucket: int, decode_wrapper: CudaGraphWrapper | None + ) -> None: + """Warm up and capture the breakable graph for ``bucket`` from the buffers.""" + for _ in range(self.num_warmup): + self._run_inner(bucket) + torch.cuda.synchronize() + stream = decode_wrapper.stream if decode_wrapper is not None else None + cap = BreakableCapture(pool=self._pool, stream=stream) + with cap: + self._outputs[bucket] = CapturedForward(*self._run_inner(bucket)) + if self._pool is None: + self._pool = cap.pool # share the pool across all subsequent buckets + cap.replay() # capture records kernels without executing; smoke-test replay + self._captures[bucket] = cap + + def _run_inner(self, num_tokens: int): + """Run the inner model over the leading ``num_tokens`` of the static buffers. + + ``num_tokens`` is the padded bucket size; the padded tail [real:bucket] is + already scrubbed to safe values (embeds=0, positions=0, + out_cache_loc=dummy_kv_slot) by :meth:`_land_input_embeds` and + ``InputBuffers.fill_input_buffers``. The embedding is NOT part of the + graph: the inner model starts from the static input-embeds buffer, so a + replay can take precomputed (e.g. merged multimodal) embeddings. + """ + ib = self.input_buffers + if self.config.model_is_mrope: + positions = ib.mrope_positions_buf[:, :num_tokens] + else: + positions = ib.positions_buf[:num_tokens] + return self.inner_model( + ib.input_ids_buf[:num_tokens], + positions, + self._ctx, + ib.out_cache_loc_buf[:num_tokens], + input_embeds=self._input_embeds_buf[:num_tokens], + ) + + def _land_input_embeds(self, embeds: torch.Tensor, bucket: int) -> None: + """Copy ``embeds`` into the static buffer's leading rows, zero the tail. + + The zeroed padded tail keeps the graphed compute over garbage-free rows + (RMSNorm of zeros is zeros; the tail is discarded by the output slice). + """ + num_tokens = embeds.shape[0] + self._input_embeds_buf[:num_tokens].copy_(embeds) + if num_tokens < bucket: + self._input_embeds_buf[num_tokens:bucket].zero_() + + def _make_dummy_batch( + self, num_tokens: int, decode_wrapper: CudaGraphWrapper | None + ) -> ForwardContext: + """Populate the static buffers + attention metadata for a dummy bs=1 extend + forward of ``num_tokens`` tokens, and return its ForwardContext. + + The prefill analogue of decode's ``_init_capture_metadata``. KV writes + go to the reserved dummy slot and the page table points at page 0, so + the forward runs (producing discarded garbage) without touching real + cache state. Backends with extra paged caches (DeepSeek-V4 DSA: SWA + + compressor + indexer state) also need per-cache block tables, or their + extend metadata comes up incomplete and the eager attention break + aborts the capture -- reuse the decode wrapper's dummy-table builder + (all zeros, the safe page 0) for those. + """ + ib = self.input_buffers + ib.input_ids_buf[:num_tokens].fill_(1) + ib.out_cache_loc_buf[:num_tokens].fill_(ib.dummy_kv_slot) + ib.positions_buf[:num_tokens].copy_( + torch.arange(num_tokens, device=self.config.device) + ) + ib.req_pool_indices_buf[:1].zero_() + ib.seq_lens_buf[:1].fill_(num_tokens) + ib.extend_seq_lens_buf[:1].fill_(num_tokens) + ib.extend_seq_lens_cpu[:1].fill_(num_tokens) + ib.extend_prefix_lens_buf[:1].zero_() + ib.extend_prefix_lens_cpu[:1].zero_() + self.req_to_page[0].zero_() # dummy request's pages -> page 0 (valid memory) + + ctx = ForwardContext( + attn_backend=self.attn_backend, + token_to_kv_pool=self.token_to_kv_pool, + req_to_page=self.req_to_page, + bs=1, + num_extends=1, + input_num_tokens=num_tokens, + forward_mode=ForwardMode.EXTEND, + capture_hidden_mode=( + CaptureHiddenMode.FULL + if self.drafter is not None + else CaptureHiddenMode.NULL + ), + ) + if self.dp_size > 1: + ctx.global_num_tokens = [num_tokens] * self.config.world_size + ctx.global_bs = [1] * self.config.world_size + extra_metadata_kwargs: dict = {} + if ( + getattr(self.attn_backend, "uses_paged_cache_groups", False) + and decode_wrapper is not None + ): + tables = decode_wrapper._capture_paged_cache_block_tables( + 1, self.token_to_kv_pool + ) + if tables is not None: + extra_metadata_kwargs["paged_cache_block_tables"] = tables + extra_metadata_kwargs["num_tokens"] = num_tokens + extra_metadata_kwargs["positions"] = ib.positions_buf[:num_tokens] + self.attn_backend.init_forward_metadata( + bs=1, + num_extends=1, + req_pool_indices=ib.req_pool_indices_buf[:1], + seq_lens=ib.seq_lens_buf[:1], + req_to_page=self.req_to_page, + forward_mode=ForwardMode.EXTEND, + extend_seq_lens=ib.extend_seq_lens_buf[:1], + extend_seq_lens_cpu=ib.extend_seq_lens_cpu[:1], + extend_prefix_lens=ib.extend_prefix_lens_buf[:1], + extend_prefix_lens_cpu=ib.extend_prefix_lens_cpu[:1], + **extra_metadata_kwargs, + ) + return ctx + + def _capture_unanimous(self, captured_ok: bool) -> bool: + """MIN-reduce capture success across the world (see ``capture``).""" + if self.config.world_group is None or self.config.world_size <= 1: + return captured_ok + from tokenspeed.runtime.distributed.process_group_manager import ( + process_group_manager as pg_manager, + ) + + cpu_group = pg_manager.get_process_group("gloo", self.config.world_group) + flag = torch.tensor([1 if captured_ok else 0], dtype=torch.int32) + torch.distributed.all_reduce( + flag, op=torch.distributed.ReduceOp.MIN, group=cpu_group + ) + unanimous = bool(flag.item()) + if not unanimous and captured_ok: + logger.warning( + "Prefill graph: a peer rank failed capture; falling back to " + "eager prefill on all ranks to keep DP/TP token counts in lockstep." + ) + return unanimous + + # ------------------------------------------------------------------ + # Replay dispatch + # ------------------------------------------------------------------ + + def can_run(self, ctx: ForwardContext, multimodal_context=None) -> bool: + """Whether this forward replays a captured graph (mirrors decode's can_run). + + A forward carrying multimodal inputs replays only when the model + exposes the embeds-only ``multimodal_input_embeds`` seam; models with + extra per-layer inputs (deepstack) run eager. + """ + if multimodal_context is not None and self._multimodal_input_embeds is None: + return False + return self._replay_bucket(ctx) is not None + + def replay( + self, + ctx: ForwardContext, + input_ids: torch.Tensor, + multimodal_context=None, + ): + """Replay the captured graph for ``ctx`` (caller checked :meth:`can_run`). + + The embedding runs eagerly here, outside the graph: a plain text + prefill gathers ``embed_tokens(input_ids)`` into the static buffer; a + multimodal prefill builds the merged text+vision embeddings via the + model's ``multimodal_input_embeds`` seam (vision encoder included) + instead -- both replay the same graphs. Then the inner stack replays + over the padded bucket and the model's eager logits tail finishes on + the real-token rows. + """ + bucket = self._replay_bucket(ctx) + assert bucket is not None, "replay() called without can_run()" + self._log_engaged_once(bucket, ctx, multimodal_context is not None) + num_tokens = ctx.input_num_tokens + input_embeds = None + if multimodal_context is not None: + input_embeds = self._multimodal_input_embeds( + input_ids, ctx, multimodal_context + ) + self._land_input_embeds( + input_embeds if input_embeds is not None else self._embed_tokens(input_ids), + bucket, + ) + with self._padded_to(ctx, bucket): + self._captures[bucket].replay() + hidden_states, aux_hidden_states = self._outputs[bucket].sliced(num_tokens) + # The eager logits tail of BaseCausalLM.forward, on the replayed hidden states. + logits_metadata = LogitsMetadata.from_forward_context(ctx) + return self.text_model.logits_processor( + input_ids, + hidden_states, + self.text_model.lm_head, + logits_metadata, + aux_hidden_states, + ) + + def _replay_bucket(self, ctx: ForwardContext) -> int | None: + """The captured bucket this forward replays, or ``None`` to run eager. + + Pure-extend AND mixed extend+decode batches are eligible: the attention + break reads the LIVE ambient ctx and dispatches the prefill/decode + split itself, while the captured token-shaped compute is uniform over + all rows (pure decode is the decode graph's job). Two ctx fields are + baked into the captured segments rather than rebound at replay -- the + draft first-step row narrowing (keyed on ``accept_lengths``) and the + ``capture_hidden_mode`` aux-hidden capture -- so a live forward carrying + different values falls back to eager rather than silently dropping the + reduce / mismatching aux. Prefix caching (cache hits and chunked-prefill + chunks 2+) IS eligible: the prefix affects only the ragged attention, + which runs entirely inside the eager break, and it adds zero new tokens, + so the padded bucket -- hence the baked EP all-to-all shape under DP -- + is identical on prefix and non-prefix ranks. + """ + if self.disable or ctx.forward_mode is None: + return None + if ctx.num_extends <= 0: + return None + if not (ctx.forward_mode.is_extend() or ctx.forward_mode.is_mixed()): + return None + if ctx.accept_lengths is not None: + return None + if ctx.capture_hidden_mode != self._captured_hidden_mode: + return None + bucket = self._select_bucket(ctx) + if bucket is None or bucket not in self._captures: + return None + return bucket + + def _select_bucket(self, ctx: ForwardContext) -> int | None: + """The padded bucket for this forward, or ``None`` to run eager. + + Under data parallelism the MoE expert-parallel all-to-all is a collective + across ALL ranks, sized from a replicated per-rank token list. The captured + graph bakes a uniform ``[bucket]*world_size`` layout, so every rank must + replay the SAME bucket or the collective desyncs (NCCL deadlock). Decide + purely from replicated global state -- the all-extend flag and the global + max token count -- so all ranks reach the identical decision/bucket with no + extra sync (mirrors the decode graph). Idle ranks run a DECODE forward, so + ``all_extend`` is False whenever any rank is idle and the graph stays off + (e.g. warmup), correctly falling back to eager. + """ + if self.dp_size <= 1 or ctx.global_num_tokens is None: + return self._padded_bucket(ctx.input_num_tokens) + if not ctx.all_extend: + return None + return self._padded_bucket(max(ctx.global_num_tokens)) + + def _padded_bucket(self, num_tokens: int) -> int | None: + """Smallest bucket >= ``num_tokens``, or ``None`` if over the largest. + + With ``--disable-cuda-graph-padding``, only an exact bucket match + replays (mirroring the decode wrapper's no-padding semantics). + """ + idx = bisect.bisect_left(self.capture_buckets, num_tokens) + if idx == len(self.capture_buckets): + return None + bucket = self.capture_buckets[idx] + if self.config.disable_cuda_graph_padding and bucket != num_tokens: + return None + return bucket + + @contextmanager + def _padded_to(self, ctx: ForwardContext, bucket: int): + """Publish ``ctx`` as the ambient live context, pinned to the padded bucket. + + The graph replays over ``bucket`` (padded) tokens; attention metadata stays + at the real count (set upstream), so the eager attention break only touches + real tokens and the padded rows produce discarded garbage. Pin + ``input_num_tokens`` to the bucket and, under DP, ``global_num_tokens`` / + ``global_bs`` to the captured uniform layout so any live read during the + break matches the baked EP shapes. The break reads ``forward_mode`` / ``bs`` + / ``num_extends`` LIVE off this same (ambient) ctx -- which we do NOT pin -- + so models split prefill vs decode and dispatch the per-mode backend + correctly with no side channel. + """ + saved = (ctx.input_num_tokens, ctx.global_num_tokens, ctx.global_bs) + ctx.input_num_tokens = bucket + if self.dp_size > 1 and ctx.global_num_tokens is not None: + ctx.global_num_tokens = [bucket] * self.config.world_size + ctx.global_bs = [1] * self.config.world_size + try: + with active_forward(ctx): + yield + finally: + ctx.input_num_tokens, ctx.global_num_tokens, ctx.global_bs = saved + + def _log_engaged_once( + self, bucket: int, ctx: ForwardContext, is_multimodal: bool + ) -> None: + kind = "multimodal" if is_multimodal else "text" + if kind in self._engaged_logged: + return + self._engaged_logged.add(kind) + logger.info( + "prefill breakable graph ENGAGED (%s): bucket=%d dp=%s mode=%s " + "(mixed prefill+decode batches supported)", + kind, + bucket, + # The replay mode actually taken (mirrors _select_bucket), a DP-debug anchor. + self.dp_size > 1 and ctx.global_num_tokens is not None, + ctx.forward_mode, + ) diff --git a/python/tokenspeed/runtime/execution/runtime_states.py b/python/tokenspeed/runtime/execution/runtime_states.py new file mode 100644 index 0000000..9c19b0a --- /dev/null +++ b/python/tokenspeed/runtime/execution/runtime_states.py @@ -0,0 +1,179 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. +"""Runtime state tensors shared by the model executor.""" + +import torch + +from tokenspeed.runtime.layers.attention.linear.mamba_state_scatter_triton import ( + fused_mamba_state_copy, + fused_mamba_state_zero, +) +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +class RuntimeStates: + """Own runtime state tensors keyed by request-pool index.""" + + def __init__( + self, + req_pool_size: int, + context_len: int, + vocab_size: int, + output_length: int, + device: str = "cuda", + mamba_pool=None, + ): + self.device = device + self.vocab_size = vocab_size + + self.valid_cache_lengths = torch.zeros( + req_pool_size + 1, dtype=torch.int32, device=device + ) + # Resolve input ids from here when overlap scheduling. + self.future_input_map = torch.empty( + (req_pool_size + 1, output_length), dtype=torch.int32, device=device + ) + self.remote_spec_candidate_ready = torch.zeros( + req_pool_size + 1, dtype=torch.bool, device=device + ) + self.linear_penalties = torch.zeros( + (req_pool_size + 1, vocab_size), dtype=torch.float32, device=device + ) + self.scaling_penalties = torch.ones( + (req_pool_size + 1, vocab_size), dtype=torch.float32, device=device + ) + self.mamba_pool = mamba_pool + + def update_valid_cache_length( + self, req_pool_indices: torch.Tensor, increment_lengths: torch.Tensor + ) -> None: + self.valid_cache_lengths.index_add_(0, req_pool_indices, increment_lengths) + + def reset_states( + self, + extend_request_pool_indices: torch.Tensor, + extend_prefix_lens: torch.Tensor, + ) -> None: + self.valid_cache_lengths[extend_request_pool_indices] = extend_prefix_lens + self.linear_penalties.index_fill_(0, extend_request_pool_indices, 0.0) + self.scaling_penalties.index_fill_(0, extend_request_pool_indices, 1.0) + self.remote_spec_candidate_ready[extend_request_pool_indices] = False + + def write_remote_spec_candidate_ids( + self, req_pool_idx: int, candidate_ids: list[int] + ) -> None: + width = self.future_input_map.shape[1] + if len(candidate_ids) != width: + raise RuntimeError( + f"remote spec candidate width mismatch: got {len(candidate_ids)}, expected {width}" + ) + ids = torch.tensor( + candidate_ids, + dtype=torch.int32, + device="cpu", + pin_memory=True, + ).to(self.device, non_blocking=True) + self.future_input_map[req_pool_idx, :width] = ids + self.remote_spec_candidate_ready[req_pool_idx] = True + + def copy_mamba_states( + self, + mamba_pool_indices: torch.Tensor, + mamba_cow_src_indices: torch.Tensor, + bs: int, + ) -> None: + """Copy Mamba states for copy-on-write requests.""" + if self.mamba_pool is None: + return + if mamba_cow_src_indices is None: + return + src_indices = mamba_cow_src_indices[:bs].long() + dst_indices = mamba_pool_indices[:bs].long() + # page_size=0 disables page-boundary filtering + fused_mamba_state_copy( + self.mamba_pool.conv_state, + src_indices, + dst_indices, + ) + fused_mamba_state_copy( + self.mamba_pool.ssm_state, + src_indices, + dst_indices, + ) + + def snapshot_mamba_checkpoints( + self, + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + cache_lengths: torch.Tensor, + page_size: int, + num_valid: int, + ) -> None: + """Copy current working Mamba states into checkpoint slots. + + src_indices/dst_indices are pre-filtered on CPU (only valid entries). + The page_size condition is checked inside the Triton kernel. + """ + if self.mamba_pool is None or num_valid == 0: + return + fused_mamba_state_copy( + self.mamba_pool.conv_state, + src_indices, + dst_indices, + cache_lengths=cache_lengths, + page_size=page_size, + ) + fused_mamba_state_copy( + self.mamba_pool.ssm_state, + src_indices, + dst_indices, + cache_lengths=cache_lengths, + page_size=page_size, + ) + + def zero_mamba_states( + self, + mamba_pool_indices: torch.Tensor, + mamba_cow_src_indices: torch.Tensor | None, + extend_prefix_lens: torch.Tensor | None, + bs: int, + ) -> None: + """Clear Mamba states for newly allocated slots without prefix state.""" + if self.mamba_pool is None: + return + pool_indices = mamba_pool_indices[:bs] + # Compute condition mask purely on GPU (no sync) + valid_pool = pool_indices != -1 + no_cow = ( + (mamba_cow_src_indices[:bs] == -1) + if mamba_cow_src_indices is not None + else torch.ones(bs, dtype=torch.bool, device=mamba_pool_indices.device) + ) + no_prefix = ( + (extend_prefix_lens[:bs] == 0) + if extend_prefix_lens is not None + else torch.ones(bs, dtype=torch.bool, device=mamba_pool_indices.device) + ) + zero_mask = valid_pool & no_cow & no_prefix + indices = torch.where(zero_mask, pool_indices, -1).long() + fused_mamba_state_zero(self.mamba_pool.conv_state, indices) + fused_mamba_state_zero(self.mamba_pool.ssm_state, indices) diff --git a/python/tokenspeed/runtime/execution/types.py b/python/tokenspeed/runtime/execution/types.py new file mode 100644 index 0000000..5aba52d --- /dev/null +++ b/python/tokenspeed/runtime/execution/types.py @@ -0,0 +1,66 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Shared result and enum types for model execution.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from tokenspeed.runtime.grammar.capturable_grammar import ( + GrammarStepCompletion, + ) + + +@dataclass +class ModelExecutionResult: + """ + Result of model execution returned to scheduler. + + This is the output from the Python executor back to the C++ scheduler. + + Attributes: + output_tokens: Sampled token IDs + output_logits: Output logits (if requested) + output_lengths: Number of tokens generated per request (for spec decoding) + """ + + output_tokens: torch.Tensor + copy_event: torch.cuda.Event | None = None + output_logits: torch.Tensor | None = None + output_lengths: torch.Tensor | None = None + grammar_completion: GrammarStepCompletion | None = None + # Per-position logprob of the sampled token, same layout as output_tokens. + # Populated unconditionally by the sampling backend so it's always + # available if any request asks for it. + output_logprobs: torch.Tensor | None = None + # Optional next-round input rows captured for PD prefill data-plane handoff. + next_input_ids: torch.Tensor | None = None + # Per-request NaN-guard flags (int32, [bs]); None when the guard is disabled. + output_nan_flags: torch.Tensor | None = None + + def sync(self) -> None: + if self.copy_event is None: + raise RuntimeError("copy_event is required before synchronizing results.") + self.copy_event.synchronize() diff --git a/python/tokenspeed/runtime/execution/weight_loader.py b/python/tokenspeed/runtime/execution/weight_loader.py new file mode 100644 index 0000000..f16a7d9 --- /dev/null +++ b/python/tokenspeed/runtime/execution/weight_loader.py @@ -0,0 +1,124 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import torch + +from tokenspeed.runtime.configs.device_config import DeviceConfig +from tokenspeed.runtime.configs.load_config import LoadConfig +from tokenspeed.runtime.configs.model_config import ModelConfig +from tokenspeed.runtime.model_loader import get_model +from tokenspeed.runtime.utils import ( + get_available_gpu_memory, + get_colorful_logger, + set_cuda_arch, +) +from tokenspeed.runtime.utils.server_args import ServerArgs +from tokenspeed.runtime.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter + +logger = get_colorful_logger(__name__) + + +class WeightLoader: + """Handles model weight loading from disk. + + This class is stateless and does not modify external state. + It returns LoadedModel with all necessary information. + """ + + @staticmethod + def load_model( + model_config: ModelConfig, + server_args: ServerArgs, + device: str, + gpu_id: int, + memory_saver_adapter: TorchMemorySaverAdapter, + ): + """Load model from disk. + + Args: + model_config: Model configuration + server_args: Server arguments + device: Device type ("cuda", "cpu") + gpu_id: GPU ID + memory_saver_adapter: Memory saver adapter + + Returns: + LoadedModel with model and dtype + """ + logger.info( + "Load weight begin. avail mem=%.2f GB", + get_available_gpu_memory(device, gpu_id), + ) + + # Reduce thread conflicts during weight loading + if device != "cpu": + torch.set_num_threads(1) + + set_cuda_arch() + + # Create load config + load_config = LoadConfig( + load_format=server_args.load_format, + download_dir=server_args.download_dir, + ext_yaml=server_args.ext_yaml, + weight_loader_prefetch_checkpoints=server_args.weight_loader_prefetch_checkpoints, + weight_loader_prefetch_num_threads=server_args.weight_loader_prefetch_num_threads, + ) + + # Load model with memory saver context. Tag as "weights" with CPU backup + # so release_memory_occupation offloads (and restores) them byte-exact. + with memory_saver_adapter.region(tag="weights", enable_cpu_backup=True): + model = get_model( + model_config=model_config, + load_config=load_config, + device_config=DeviceConfig(device), + ) + + # Load KV cache scaling factors if using FP8 + if server_args.kv_cache_dtype == "fp8_e4m3": + if server_args.quantization_param_path is not None: + if callable(getattr(model, "load_kv_cache_scales", None)): + model.load_kv_cache_scales(server_args.quantization_param_path) + logger.info( + "Loaded KV cache scaling factors from %s", + server_args.quantization_param_path, + ) + else: + raise RuntimeError( + "Using FP8 KV cache and scaling factors provided but " + f"model {model.__class__} does not support loading scaling factors." + ) + else: + logger.warning( + "Using FP8 KV cache but no scaling factors provided. " + "Defaulting to scaling factors of 1.0. " + "This may lead to less accurate results!" + ) + + dtype = model_config.dtype + + logger.info( + "Load weight end. type=%s, dtype=%s, avail mem=%.2f GB", + type(model).__name__, + dtype, + get_available_gpu_memory(device, gpu_id), + ) + + return model diff --git a/python/tokenspeed/runtime/grammar/__init__.py b/python/tokenspeed/runtime/grammar/__init__.py new file mode 100644 index 0000000..f34a629 --- /dev/null +++ b/python/tokenspeed/runtime/grammar/__init__.py @@ -0,0 +1,21 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Structured-output and grammar-guided decoding helpers.""" diff --git a/python/tokenspeed/runtime/grammar/base_grammar_backend.py b/python/tokenspeed/runtime/grammar/base_grammar_backend.py new file mode 100755 index 0000000..6de3f9d --- /dev/null +++ b/python/tokenspeed/runtime/grammar/base_grammar_backend.py @@ -0,0 +1,457 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright 2023-2024 SGLang Team +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Base classes for grammar-guided constrained decoding backends.""" + +import time +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass +from threading import Event, Lock +from typing import Any + +from tokenspeed.runtime.utils.server_args import ServerArgs + + +@dataclass +class CacheEntry: + + value: Any + event: Event + + # Set non-None when a worker begins compiling for this key. Doubles as + # an ownership marker so a second ``init_value`` call for the same key + # waits on ``event`` instead of starting a redundant compile. + started_at: float | None = None + + # Shared future for the in-flight compile, set by + # ``get_cached_or_future_value`` under the cache lock. Concurrent + # callers for the same key reuse this instead of submitting duplicate + # tasks that would saturate the executor on repeated same-schema traffic. + future: Future | None = None + + +class BaseGrammarObject: + """Base for compiled grammar wrappers. Concrete backends subclass this.""" + + is_invalid: bool = False + expires_at: float | None = None + + @property + def expired(self) -> bool: + + return False + + +class InvalidGrammarObject(BaseGrammarObject): + """Sentinel cached when a grammar fails to compile (or compile times out). + + Carries the underlying error string so the request abort message can + surface the actual reason — bad regex, malformed schema, timeout, etc. + + ``expires_at`` (monotonic seconds) lets transient failures decay so a + one-off slow compile doesn't poison a valid key forever. ``None`` means + the failure is permanent (e.g. malformed schema — recompiling won't help). + """ + + is_invalid = True + + def __init__( + self, + error_message: str = "Unknown grammar error", + expires_at: float | None = None, + ) -> None: + + super().__init__() + + self.error_message = error_message + self.expires_at = expires_at + + @property + def expired(self) -> bool: + + return self.expires_at is not None and time.monotonic() >= self.expires_at + + def __repr__(self) -> str: + + return ( + f"InvalidGrammarObject(error_message={self.error_message!r}, " + f"expires_at={self.expires_at!r})" + ) + + +@dataclass +class TimeoutHistory: + """Per-key bookkeeping so the cache can escalate to a permanent failure + after enough transient timeouts (a single slow compile shouldn't poison + the key, but a key that *consistently* times out is broken). + """ + + count: int = 0 + + +class BaseGrammarBackend: + """Base backend with shared cache management for grammar objects.""" + + def __init__(self): + + self.executor = ThreadPoolExecutor() + + self.cache: dict[tuple[str, str], CacheEntry] = {} + self.cache_lock = Lock() + + # Tracks transient-timeout attempts per key. Cleared on a successful + # compile (in init_value) or on reset(). + self.timeout_history: dict[tuple[str, str], TimeoutHistory] = {} + + # ------------------------------------------------------------------ + # Subclass hook + # ------------------------------------------------------------------ + + def init_value_impl( + self, key: tuple[str, str], require_reasoning: bool + ) -> BaseGrammarObject: + raise NotImplementedError() + + # ------------------------------------------------------------------ + # Cache API + # ------------------------------------------------------------------ + + def init_value(self, key: tuple[str, str]) -> BaseGrammarObject: + + with self.cache_lock: + + entry = self.cache.get(key) + + if entry is None: + + entry = CacheEntry(None, Event(), started_at=time.monotonic()) + self.cache[key] = entry + + cache_hit = False + + elif entry.event.is_set() or entry.started_at is not None: + + # Either already compiled, or another worker claimed the + # compile (``started_at`` set). Wait on the event below. + cache_hit = True + + else: + + # Entry was pre-created by ``get_cached_or_future_value`` + # but no worker has claimed it yet — take ownership. + entry.started_at = time.monotonic() + cache_hit = False + + if cache_hit: + + entry.event.wait() + + else: + + # Backends should return InvalidGrammarObject(message); accept + # legacy None as "unknown error" for safety. + if (value := self.init_value_impl(key, False)) is None: + + value = InvalidGrammarObject() + + with self.cache_lock: + + is_active = self.cache.get(key) is entry + + if is_active: + + # If a parallel caller (e.g. GrammarManager on + # compile timeout) wrote an InvalidGrammarObject + # while we were compiling, keep that marker as long + # as it hasn't expired — otherwise a slow compile + # finishing a moment after timeout would silently + # un-invalidate the cache. Once the marker expires + # we let the freshly-compiled value take over so a + # transient slow compile doesn't poison the key + # forever. + cached = entry.value + + if not ( + cached is not None and cached.is_invalid and not cached.expired + ): + entry.value = value + + if not value.is_invalid: + # Successful compile — clear timeout bookkeeping. + self.timeout_history.pop(key, None) + + else: + + # Orphan: our entry was evicted (e.g. its timeout + # marker expired and ``get_cached_or_future_value`` + # replaced it). Don't touch shared state — + # ``timeout_history`` now belongs to the new attempt + # — but still publish our compiled value on the old + # entry so any future/wait still holding a reference + # to it resolves with a usable result instead of + # deadlocking or returning None. + entry.value = value + + entry.event.set() + + if entry.value.is_invalid: + + return entry.value + + return entry.value.copy() + + def get_cached_or_future_value( + self, key: tuple[str, str] + ) -> tuple[BaseGrammarObject | Future, bool]: + """Return (value, cache_hit). + + On cache hit: value is either a fresh grammar copy or an + InvalidGrammarObject carrying the original compile error. + On miss: value is a Future that resolves to the same. + + Expired InvalidGrammarObject markers are evicted on lookup so a + transient timeout decays into a retry instead of poisoning the key. + """ + with self.cache_lock: + + entry = self.cache.get(key) + + if ( + entry is not None + and entry.event.is_set() + and entry.value.is_invalid + and entry.value.expired + ): + # Drop the stale marker. + del self.cache[key] + entry = None + + if entry is not None and entry.event.is_set(): + + if entry.value.is_invalid: + + return entry.value, True + + return entry.value.copy(), True + + # In-flight compile — share its future so concurrent callers + # for the same key don't each submit duplicate work that + # would park executor workers on the same event. + if entry is not None and entry.future is not None: + + return entry.future, False + + # Pre-create the entry under the lock so any caller that + # races in after us finds our shared future instead of + # submitting its own. + if entry is None: + + entry = CacheEntry(value=None, event=Event()) + self.cache[key] = entry + + entry.future = self.executor.submit(self.init_value, key) + + return entry.future, False + + def compile_started_at(self, key: tuple[str, str]) -> float | None: + """Monotonic timestamp at which the worker thread began running + ``init_value_impl`` for ``key``, or None if no compile has ever + started for this key (no cache entry, or the entry was created + directly via ``cache_invalid`` / ``record_compile_timeout`` without + a backing compile). + + Returned even after the compile has finished — callers want + "compile-only elapsed", which is ``now - started_at`` regardless of + completion. Otherwise a worker that finishes between a caller's + ``future.done()`` check and the elapsed-time check would silently + flip the elapsed calculation back to wall-clock-from-submit and + could trigger a spurious timeout against a request that actually + succeeded. + """ + with self.cache_lock: + + entry = self.cache.get(key) + + return entry.started_at if entry is not None else None + + def record_compile_timeout( + self, + key: tuple[str, str], + error_message: str, + ttl_secs: float, + max_retries: int, + ) -> None: + """Cache a compile-timeout marker for ``key``. + + Each call increments a per-key attempt counter. While the count is at + or below ``max_retries`` the marker has a finite TTL so the next + request (after the marker expires) gets a fresh compile attempt. + Once the count crosses the threshold the marker is escalated to + permanent (no TTL) — at that point the compiler is consistently + broken for this key and there's nothing to gain by retrying. + + Spurious-timeout safety: if the cache already holds a valid grammar + (the worker just finished compiling, racing this call), this is a + no-op — we don't penalize a working key, and we don't clobber a + freshly-committed result with a stale timeout. + + Counter reset semantics: + - cleared on a successful compile (see ``init_value``) + - cleared on ``reset()`` + """ + with self.cache_lock: + + entry = self.cache.get(key) + + if ( + entry is not None + and entry.event.is_set() + and not entry.value.is_invalid + ): + # A valid grammar landed in the cache while we were about to + # declare timeout. Drop the timeout — the grammar works. + return + + history = self.timeout_history.setdefault(key, TimeoutHistory()) + history.count += 1 + + if history.count > max_retries: + + expires_at = None # permanent + + error_message = ( + f"{error_message} (gave up after {history.count - 1} retries)" + ) + + else: + + expires_at = time.monotonic() + ttl_secs + + invalid = InvalidGrammarObject(error_message, expires_at=expires_at) + + if entry is None: + + event = Event() + event.set() + + self.cache[key] = CacheEntry(invalid, event) + + elif not entry.event.is_set(): + + entry.value = invalid + entry.event.set() + + else: + + # entry.value.is_invalid (valid was filtered above) — refresh. + entry.value = invalid + + def cache_invalid(self, key: tuple[str, str], error_message: str) -> None: + """Cache a permanent compile failure (e.g. bad syntax — retrying + won't help). Use ``record_compile_timeout`` for transient failures.""" + invalid = InvalidGrammarObject(error_message, expires_at=None) + + with self.cache_lock: + + entry = self.cache.get(key) + + if entry is None: + + event = Event() + event.set() + + self.cache[key] = CacheEntry(invalid, event) + + elif not entry.event.is_set(): + + entry.value = invalid + entry.event.set() + + elif entry.value.is_invalid: + + entry.value = invalid + + def reset(self): + + with self.cache_lock: + + self.cache.clear() + self.timeout_history.clear() + + +def create_grammar_backend(server_args: ServerArgs, tokenizer, vocab_size): + + if server_args.grammar_backend == "none": + + return None + + elif server_args.grammar_backend == "xgrammar": + + from tokenspeed.runtime.grammar.xgrammar_backend import ( + XGrammarGrammarBackend, + ) + + grammar_backend = XGrammarGrammarBackend( + tokenizer, + vocab_size=vocab_size, + disable_any_whitespace=server_args.disable_any_whitespace, + ) + + else: + + raise ValueError(f"Invalid grammar backend: {server_args.grammar_backend}") + + # Reasoning + grammar deferral lives in the OpenAI serving layer now: + # response_format is rewritten into an xgrammar structural_tag whose + # trigger covers the post-reasoning preamble. gpt-oss is wired today + # (``<|start|>assistant<|channel|>final<|message|>``); other reasoning + # parsers (qwen3-thinking, deepseek-r1, ...) can be added the same + # way as needed. The previous token-id-based ``ReasonerGrammarBackend`` + # wrapper has been removed: it couldn't handle multi-token channel + # preambles and carried known P1 bugs around state cloning + reasoning + # initialization. + return grammar_backend + + +def get_apply_vocab_mask_func(grammar_backend: str): + """Return the backend-specific in-place vocab-mask-apply function. + + The function's signature is ``(logits, vocab_mask) -> None``. It is + stored on ``SamplingBatchInfo.apply_vocab_mask`` so the captured + sampler can call it without branching on backend. Only xgrammar is + wired up today; add branches here when new backends are added. + """ + if grammar_backend == "xgrammar": + + from xgrammar import apply_token_bitmask_inplace + + # xgrammar's native signature is (logits, bitmask, *, ...); + # adapt to the canonical (logits, vocab_mask) kwargs the + # sampler uses. + def _apply(logits, vocab_mask): + apply_token_bitmask_inplace(logits, vocab_mask) + + return _apply + + raise ValueError(f"Unsupported grammar backend: {grammar_backend}") diff --git a/python/tokenspeed/runtime/grammar/capturable_grammar.py b/python/tokenspeed/runtime/grammar/capturable_grammar.py new file mode 100644 index 0000000..32c383c --- /dev/null +++ b/python/tokenspeed/runtime/grammar/capturable_grammar.py @@ -0,0 +1,605 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Host-function-driven grammar pipeline. + +Captures ``fill_next_token_bitmask`` + H2D as graph nodes via +``cudaLaunchHostFunc`` (see ``tokenspeed.runtime.utils.hostfunc``) so the +grammar fill lives inside the CUDA graph and overlaps with the model +forward on a side stream. + +Deferred advance: the matcher advance for step N's sampled tokens runs +inside step N+1's ``build`` hostfunc, not at the tail of step N's +graph. ``schedule_post_sampler`` does a main-stream D2H of the sampler +output into a pinned buffer shared across steps. The next step's +``fork_event`` (recorded on main inside captured ``schedule_fill``) +transitively waits for this D2H before the side-stream build reads +the pinned memory. Main only joins on ``bitmask_event`` before +apply_mask, so forward(N+1) overlaps with the prev-step matcher +advance and this step's mask fill. The shared pinned buffer is +read-only for the next step's build; ``post_process`` reads its own +per-step CPU tensor produced by ``.to('cpu')`` in execute_forward_op, +so step N+1 overwriting pinned does not race with commit(N). +""" + +from __future__ import annotations + +import queue +import threading +from dataclasses import dataclass, field + +import torch + +from tokenspeed.runtime.utils import get_colorful_logger +from tokenspeed.runtime.utils.hostfunc import hostfunc + +logger = get_colorful_logger(__name__) + + +@dataclass +class GrammarStepInputs: + """Per-batch grammar state assembled by the event loop. + + ``grammars[i]`` is the matcher for request ``i`` in the batch, or + ``None`` if that request has no grammar. ``advance_mask[i] == False`` + means slot ``i`` is an intermediate chunked-prefill chunk whose + sampled token must NOT advance the matcher (the output is discarded + by post_process). ``advance_mask is None`` means "advance all". + + Pass ``None`` instead of an instance when no request in the batch + has a grammar — the executor's setup_grammar_step short-circuits. + """ + + grammars: list[object | None] + advance_mask: list[bool] | None = None + + +@dataclass +class GrammarStepCompletion: + """Per-step handoff between the build hostfunc and post_process. + + The next step's ``build`` normally sets ``event`` after advancing + the matcher and recording ``terminated_at[i]`` — the index of the + token (within this step's accepted-token chain) that terminated + request ``i``'s matcher, or -1. If no next step is dispatched + (request was the last live one), post_process advances on the + host instead; ``lock`` ensures exactly one path wins. + + ``advance_mask[i] == False`` means slot i's sampled token was garbage + (intermediate chunked-prefill chunk) and MUST NOT advance the matcher. + """ + + event: threading.Event = field(default_factory=threading.Event) + terminated_at: list[int] = field(default_factory=list) + + lock: threading.Lock = field(default_factory=threading.Lock) + grammars: list | None = None + + bs: int = 0 + tokens_per_req: int = 1 + + advance_mask: list[bool] | None = None + + +class CapturableGrammarExecutor: + """Buffers + hostfuncs for an in-graph grammar fill + apply.""" + + def __init__( + self, + max_bs: int, + vocab_size: int, + max_tokens_per_req: int = 1, + device: torch.device = torch.device("cuda"), + ) -> None: + + self.max_bs = max_bs + self.vocab_size = vocab_size + self.max_tokens_per_req = max(1, max_tokens_per_req) + self.bitmask_width = (vocab_size + 31) // 32 + self.device = device + + n_rows = max_bs * self.max_tokens_per_req + + with torch.device(device): + + self.bitmask = torch.full( + (n_rows, self.bitmask_width), -1, dtype=torch.int32 + ) + + self.bitmask_host = torch.full( + (n_rows, self.bitmask_width), + -1, + dtype=torch.int32, + pin_memory=True, + ) + + # Draft tokens for spec verify; col 0 is unused for non-spec. + self.candidates_host = torch.zeros( + (max_bs, self.max_tokens_per_req), dtype=torch.int32, pin_memory=True + ) + + # Sampler output copied to pinned memory at the tail of each step. + # Read by the NEXT step's build (via fork_event ordering); overwritten + # at the next step's tail (strictly after that step's bitmask_event, + # hence after build's read). + self.output_tokens_host = torch.zeros( + (max_bs * self.max_tokens_per_req,), dtype=torch.int32, pin_memory=True + ) + + self.accept_lengths_host = torch.zeros( + (max_bs,), dtype=torch.int32, pin_memory=True + ) + + # One queue entry per replay: CPU pushes via add_batch, the + # fetch_batch hostfunc pops + shifts current into prev. + self.queue = queue.Queue() + self.current_batch: dict | None = None + self.prev_batch: dict | None = None + + self.stream = torch.cuda.Stream() + self.fork_event = torch.cuda.Event() + self.bitmask_event = torch.cuda.Event() + + def add_batch( + self, + grammars: list, + bs: int, + has_candidates: bool = False, + tokens_per_req: int = 1, + advance_mask: list[bool] | None = None, + ) -> GrammarStepCompletion: + """Push request state for the next captured iteration. + + Must be called exactly once before each replay (including + capture warmup); ``fetch_batch`` raises on an empty queue. + + ``advance_mask[i] == False`` disables matcher advance for slot + i at this step — set it for intermediate chunked-prefill chunks + whose sampled tokens are discarded by output_processor. + ``None`` defaults to all-True. + """ + grammars_list = list(grammars) + + if advance_mask is None: + + advance_list = [True] * bs + + else: + + advance_list = list(advance_mask) + + if len(advance_list) != bs: + raise ValueError(f"advance_mask length {len(advance_list)} != bs {bs}") + + completion = GrammarStepCompletion( + grammars=grammars_list, + bs=bs, + tokens_per_req=tokens_per_req, + advance_mask=advance_list, + ) + + self.queue.put( + { + "grammars": grammars_list, + "bs": bs, + "has_candidates": has_candidates, + "tokens_per_req": tokens_per_req, + "advance_mask": advance_list, + "completion": completion, + } + ) + + return completion + + def reset_state(self) -> None: + """Drop any warmup-run state held by prev/current pointers.""" + self.prev_batch = None + self.current_batch = None + + @hostfunc + def fetch_batch(self) -> None: + """Pop batch from queue, setting prev to current and current to the new batch.""" + self.prev_batch = self.current_batch + self.current_batch = self.queue.get_nowait() + + @hostfunc + def build(self) -> None: + """Advance matcher by prev step's outputs, then fill this step's bitmask.""" + self.prev_batch and self._advance_prev(self.prev_batch) + self._fill_current(self.current_batch) + + def _advance_prev(self, prev: dict) -> None: + """Advance each prev-step grammar by its accepted tokens and record + which (if any) terminated in this step.""" + completion: GrammarStepCompletion = prev["completion"] + + # Lock serializes with post_process's host-side fallback: + # exactly one path advances the matcher and fires the event. + with completion.lock: + + if completion.event.is_set(): + return + + grammars = prev["grammars"] + stride = prev["tokens_per_req"] + bs = prev["bs"] + advance_mask = prev["advance_mask"] + terminated_at = [-1] * bs + + for i, grammar in enumerate(grammars): + + if ( + grammar is None + or grammar.finished + or grammar.is_terminated() + or not advance_mask[i] + ): + continue + + n_accepted = int(self.accept_lengths_host[i].item()) + + for j in range(n_accepted): + + tok = int(self.output_tokens_host[i * stride + j].item()) + + try: + + grammar.accept_token(tok) + + except Exception: + + break + + if grammar.is_terminated(): + + terminated_at[i] = j + break + + completion.terminated_at = terminated_at + completion.event.set() + + def _fill_current(self, batch: dict | None) -> None: + """Fill bitmask_host for this step's grammars.""" + if batch is None: + + self.bitmask_host.fill_(-1) + + return + + grammars = batch["grammars"] + bs = batch["bs"] + n = self.max_tokens_per_req + has_candidates = batch["has_candidates"] + + # Spec verify binds bitmask[:bs*n] for rejection_sampling; + # non-spec binds bitmask[:bs] for Sampler.sample. + per_req_rows = n if has_candidates else 1 + self.bitmask_host[: bs * n].fill_(-1) + + for i, grammar in enumerate(grammars): + + if grammar is None or grammar.finished or grammar.is_terminated(): + + continue + + row_base = i * per_req_rows + advanced = 0 + + for pos in range(n): + + if grammar.is_terminated(): + + break + + grammar.fill_vocab_mask(self.bitmask_host, row_base + pos) + + if pos + 1 == n or not has_candidates: + + break + + # col 0 was consumed by a previous step's advance; + # walk cols 1..n-1 to produce per-position masks. + next_tok = int(self.candidates_host[i, pos + 1].item()) + + if not grammar.try_accept_token(next_tok): + + break + + advanced += 1 + + # Undo the draft walk — the real advance happens in the + # NEXT step's build based on the sampler's accepted count. + if advanced: + + grammar.rollback(advanced) + + def schedule_fill( + self, + input_ids_buf_slice: torch.Tensor | None = None, + ) -> None: + """Fork grammar work onto the side stream for this step. + + Side stream: wait(fork_event) → D2H candidates (spec) → + fetch_batch → build → H2D bitmask → bitmask_event. Main rejoins + via wait_bitmask before apply_mask; forward on main overlaps + with the advance + fill on the side stream. + """ + self.fork_event.record() + + with torch.cuda.stream(self.stream): + + torch.cuda.current_stream().wait_event(self.fork_event) + + if input_ids_buf_slice is not None: + + bs = input_ids_buf_slice.shape[0] // self.max_tokens_per_req + + self.candidates_host[:bs].copy_( + input_ids_buf_slice.view(bs, self.max_tokens_per_req), + non_blocking=True, + ) + + self.fetch_batch() + self.build() + + self.bitmask.copy_(self.bitmask_host, non_blocking=True) + self.bitmask_event.record() + + def wait_bitmask(self) -> None: + """Join the side stream on the main stream before apply_mask.""" + torch.cuda.current_stream().wait_event(self.bitmask_event) + + def schedule_post_sampler( + self, + output_tokens: torch.Tensor, + accept_lengths: torch.Tensor, + ) -> None: + """Main-stream D2H of sampler output into the pinned buffer.""" + + n = output_tokens.numel() + self.output_tokens_host[:n].copy_(output_tokens.flatten(), non_blocking=True) + + m = accept_lengths.numel() + self.accept_lengths_host[:m].copy_(accept_lengths, non_blocking=True) + + +class EagerGrammarBuffers: + """GPU + pinned-CPU buffers for the non-CUDA grammar fallback path. + + ``CapturableGrammarExecutor`` uses ``cudaLaunchHostFunc`` for its + side-stream fill, which is CUDA-only. On HIP / CPU we fall back to a + synchronous D2H + CPU xgrammar fill + H2D, which needs its own + pre-allocated buffers (kept off ``InputBuffers`` so model-input state + isn't muddled with grammar state). + """ + + def __init__( + self, + max_bs: int, + vocab_size: int, + max_tokens_per_req: int = 1, + device: str = "cuda", + ) -> None: + + self.max_bs = max_bs + self.vocab_bitmask_width = (vocab_size + 31) // 32 + self.max_tokens_per_req = max(1, max_tokens_per_req) + + with torch.device(device): + self.vocab_mask_buf = torch.full( + (max_bs, self.vocab_bitmask_width), -1, dtype=torch.int32 + ) + # Spec-verify grammar bitmask: flat [max_bs * n, width] to match + # the apply kernel's expected shape. + if self.max_tokens_per_req > 1: + self.vocab_mask_spec_buf = torch.full( + (max_bs * self.max_tokens_per_req, self.vocab_bitmask_width), + -1, + dtype=torch.int32, + ) + + # Pinned staging for the H2D copy. + self.vocab_mask_cpu_buf = torch.full( + (max_bs, self.vocab_bitmask_width), + -1, + dtype=torch.int32, + pin_memory=True, + ) + if self.max_tokens_per_req > 1: + self.vocab_mask_spec_cpu_buf = torch.full( + (max_bs * self.max_tokens_per_req, self.vocab_bitmask_width), + -1, + dtype=torch.int32, + pin_memory=True, + ) + # Draft candidates D2H'd per step so the CPU grammar fill can + # walk the draft chain position-by-position. + self.candidates_cpu_buf = torch.zeros( + (max_bs, self.max_tokens_per_req), + dtype=torch.int32, + pin_memory=True, + ) + + +def bind_grammar_mask_buf( + info, + eager_buffers: EagerGrammarBuffers | None, + bs: int, + *, + spec: bool, + capturable: CapturableGrammarExecutor | None, + grammar_backend: str, +) -> None: + """Bind the preallocated grammar bitmask buffer onto ``info``. + + The captured sampler always takes the apply-mask branch when a buffer + is bound; for non-grammar batches the buffer stays all-ones so apply + is a no-op. When no buffer is allocated (grammar disabled) this is a + no-op and sampling skips the mask entirely. + """ + if capturable is None and eager_buffers is None: + return + + from tokenspeed.runtime.grammar.base_grammar_backend import ( + get_apply_vocab_mask_func, + ) + + if capturable is not None: + n = capturable.max_tokens_per_req + info.vocab_mask = ( + capturable.bitmask[: bs * n] if spec else capturable.bitmask[:bs] + ) + elif spec and eager_buffers.max_tokens_per_req > 1: + info.vocab_mask = eager_buffers.vocab_mask_spec_buf[ + : bs * eager_buffers.max_tokens_per_req + ] + else: + info.vocab_mask = eager_buffers.vocab_mask_buf[:bs] + info.apply_vocab_mask = get_apply_vocab_mask_func(grammar_backend) + + +def _fill_eager_bitmask( + grammars: list, + bs: int, + eager_buffers: EagerGrammarBuffers, + spec_num_tokens: int, + is_spec_decode: bool, + input_ids_buf, +) -> None: + """Sync, walk grammars on host, H2D the bitmask. Non-CUDA path only.""" + if is_spec_decode: + eager_buffers.candidates_cpu_buf[:bs].copy_( + input_ids_buf[: bs * spec_num_tokens].view(bs, spec_num_tokens), + non_blocking=True, + ) + sync_ev = torch.cuda.Event() + sync_ev.record() + sync_ev.synchronize() + cand_cpu = eager_buffers.candidates_cpu_buf + active = bs * spec_num_tokens + cpu_buf = eager_buffers.vocab_mask_spec_cpu_buf + gpu_buf = eager_buffers.vocab_mask_spec_buf + cpu_buf[:active].fill_(-1) + for i, grammar in enumerate(grammars): + if grammar is None or grammar.finished or grammar.is_terminated(): + continue + row_base = i * spec_num_tokens + advanced = 0 + for pos in range(spec_num_tokens): + if grammar.is_terminated(): + break + grammar.fill_vocab_mask(cpu_buf, row_base + pos) + if pos + 1 == spec_num_tokens: + break + next_tok = int(cand_cpu[i, pos + 1].item()) + if not grammar.try_accept_token(next_tok): + break + advanced += 1 + if advanced: + grammar.rollback(advanced) + gpu_buf[:active].copy_(cpu_buf[:active], non_blocking=True) + else: + cpu_buf = eager_buffers.vocab_mask_cpu_buf + gpu_buf = eager_buffers.vocab_mask_buf + cpu_buf[:bs].fill_(-1) + for i, grammar in enumerate(grammars): + if grammar and not grammar.finished and not grammar.is_terminated(): + grammar.fill_vocab_mask(cpu_buf, i) + gpu_buf[:bs].copy_(cpu_buf[:bs], non_blocking=True) + + +GrammarRuntime = CapturableGrammarExecutor | EagerGrammarBuffers + + +def setup_grammar_step( + *, + sampling_info, + bs: int, + is_spec_decode: bool, + spec_num_tokens: int, + grammar_inputs: GrammarStepInputs | None, + grammar_runtime: GrammarRuntime | None, + input_ids_buf, + grammar_backend: str, +) -> GrammarStepCompletion | None: + """Bind the bitmask buffer and dispatch one step of grammar work. + + ``grammar_runtime`` is one of: + - ``CapturableGrammarExecutor`` (CUDA): enqueues an ``add_batch`` so + the side-stream hostfunc fills the bitmask in parallel with the + forward. Returns the per-step ``GrammarStepCompletion``. + - ``EagerGrammarBuffers`` (non-CUDA fallback): syncs, runs the + xgrammar fill on host, H2Ds the bitmask. Returns ``None``. + - ``None``: grammar disabled, no-op. + """ + if grammar_runtime is None: + return None + + capturable = ( + grammar_runtime + if isinstance(grammar_runtime, CapturableGrammarExecutor) + else None + ) + eager_buffers = ( + grammar_runtime if isinstance(grammar_runtime, EagerGrammarBuffers) else None + ) + + bind_grammar_mask_buf( + sampling_info, + eager_buffers, + bs, + spec=is_spec_decode, + capturable=capturable, + grammar_backend=grammar_backend, + ) + + grammars = grammar_inputs.grammars if grammar_inputs is not None else [None] * bs + advance_mask = grammar_inputs.advance_mask if grammar_inputs is not None else None + + if capturable is not None: + # Always push (even all-None) to keep the captured hostfunc queue + # 1:1 with replays. + tokens_per_req = spec_num_tokens if is_spec_decode else 1 + return capturable.add_batch( + grammars=grammars, + bs=bs, + has_candidates=is_spec_decode, + tokens_per_req=tokens_per_req, + advance_mask=advance_mask, + ) + + # Fill the bound buffer every step. When no request has a grammar we + # still need to write all-ones (-1) to clear any leftover bits from a + # prior grammar-batch — the captured graph reads from this same memory + # whether or not we filled it this step. + if any(grammars): + _fill_eager_bitmask( + grammars, + bs, + eager_buffers, + spec_num_tokens, + is_spec_decode, + input_ids_buf, + ) + elif is_spec_decode and eager_buffers.max_tokens_per_req > 1: + eager_buffers.vocab_mask_spec_buf[: bs * spec_num_tokens].fill_(-1) + else: + eager_buffers.vocab_mask_buf[:bs].fill_(-1) + return None diff --git a/python/tokenspeed/runtime/grammar/grammar_manager.py b/python/tokenspeed/runtime/grammar/grammar_manager.py new file mode 100644 index 0000000..e756669 --- /dev/null +++ b/python/tokenspeed/runtime/grammar/grammar_manager.py @@ -0,0 +1,370 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright 2023-2024 SGLang Team +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import time +from concurrent import futures +from typing import TYPE_CHECKING + +import torch + +from tokenspeed.runtime.distributed.process_group_manager import ( + process_group_manager as pg_manager, +) +from tokenspeed.runtime.grammar.base_grammar_backend import ( + InvalidGrammarObject, + create_grammar_backend, +) +from tokenspeed.runtime.utils import get_colorful_logger + +if TYPE_CHECKING: + from tokenspeed.runtime.engine.generation_output_processor import RequestState + from tokenspeed.runtime.utils.server_args import ServerArgs + +logger = get_colorful_logger(__name__) + + +# A pending grammar request: the spec must be submitted to the C++ scheduler and +# the state registered with the executor *only* after its grammar future resolves. +QueueEntry = tuple["object", "RequestState", "object"] # (spec, state, bootstrap_info) + + +class GrammarManager: + def __init__( + self, + server_args: ServerArgs, + tokenizer, + vocab_size: int, + ) -> None: + + self.server_args = server_args + self.grammar_queue: list[QueueEntry] = [] + + self.compile_timeout_secs = float(server_args.grammar_compile_timeout_secs) + self.compile_max_retries = int(server_args.grammar_compile_max_retries) + + # Backend is None when (a) the user disabled grammar via + # --grammar-backend none or (b) tokenizer init is skipped (no tokenizer + # → can't build a backend). Either way, requests with grammar fields + # get rejected at admission time in process_req_with_grammar. + if server_args.skip_tokenizer_init: + self.grammar_backend = None + + else: + self.grammar_backend = create_grammar_backend( + server_args, tokenizer, vocab_size + ) + + # Grammar admission must be coherent across attention TP ranks: every rank + # in the group sees the same recv_reqs (broadcast in RequestHandler.recv_reqs), + # so they must agree on which requests have a ready grammar before admitting + # any of them. Different DP ranks see different requests and don't sync. + attn_group = server_args.mapping.attn.tp_group + + if len(attn_group) > 1 and pg_manager.has_process_group("gloo", attn_group): + self.grammar_sync_group = pg_manager.get_process_group("gloo", attn_group) + self.grammar_sync_size = len(attn_group) + + else: + self.grammar_sync_group = None + self.grammar_sync_size = 1 + + def __len__(self): + + return len(self.grammar_queue) + + def clear(self): + + if self.grammar_backend: + self.grammar_backend.reset() + + def mark_abort(self, rid: str) -> None: + """Cancel a queued request whose grammar is still compiling. + + Paired with ``OutputProcesser.mark_abort`` in the event loop: + together they cover both already-registered requests and those + still blocked on a grammar future. Without this, an aborted + request would finish compiling, get admitted, and only then be + noticed as aborted — burning capacity in the meantime. + """ + for spec, state, _ in self.grammar_queue: + if spec.request_id == rid: + logger.debug("Abort grammar queue request. rid=%s", rid) + + # Don't cancel the compile future: it's shared across + # every concurrent request for the same grammar key + # (see ``get_cached_or_future_value``), so cancelling + # would raise CancelledError on still-valid waiters. + # The compile runs to completion in the background and + # its result lands in the cache for future reuse. + state.set_finish_with_abort("Aborted by AbortReq.") + + # The queue entry is cleaned up by the next + # get_ready_grammar_requests pass (which treats + # state.finished as "ready, but don't admit"). + return + + def process_req_with_grammar(self, state: RequestState) -> bool: + """Attach grammar (or future) to ``state``. + + Returns True if the request is admittable now (no grammar, cache hit, or + already aborted), False if it must be queued until its future resolves. + """ + sp = state.sampling_params + + if ( + sp.json_schema is None + and sp.regex is None + and sp.ebnf is None + and sp.structural_tag is None + ): + return True + + if self.grammar_backend is None: + state.set_finish_with_abort( + "Grammar-based generation (json_schema, regex, ebnf, structural_tag) " + "is not supported when the server is launched with --grammar-backend none" + ) + + return True + + if sp.json_schema is not None: + key = ("json", sp.json_schema) + + elif sp.regex is not None: + key = ("regex", sp.regex) + + elif sp.ebnf is not None: + key = ("ebnf", sp.ebnf) + + else: + key = ("structural_tag", sp.structural_tag) + + value, cache_hit = self.grammar_backend.get_cached_or_future_value(key) + state.grammar_key = key + + if cache_hit: + if value.is_invalid: + state.set_finish_with_abort( + f"Failed to compile {key[0]} grammar: {value.error_message}" + ) + + state.grammar = None + + else: + state.grammar = value + + return True + + # Compile is in flight; caller should add to queue via add_to_queue. + state.grammar = value # Future + return False + + def add_to_queue(self, spec, state: RequestState, bootstrap_info) -> None: + # Per-state queue timestamp bounds the time a request can spend in + # the executor's pending queue before it's picked up by a worker. + # The compile-only budget (after the worker starts) is measured + # against backend.compile_started_at(key) instead. + state.grammar_queued_ts = time.monotonic() + self.grammar_queue.append((spec, state, bootstrap_info)) + + def get_ready_grammar_requests(self) -> list[QueueEntry]: + """Promote queued requests whose grammar is ready (or has timed out). + + Per-rank: scan futures, mark ready/failed locally. + Cross-rank (attn TP > 1): all_gather indices and admit only the + intersection of ready sets and the union of failed sets, so every rank + admits the same requests in the same iteration. + + Caller is responsible for invoking this every loop iteration when + ``grammar_sync_size > 1`` so the collective stays in sync; with size 1 + it's a no-op when the queue is empty. + """ + # Queue length is identical on every attn-TP rank: recv_reqs + # broadcasts the new/abort sets, so add_to_queue and mark_abort + # run in lockstep, and pops below only happen after the cross- + # rank consensus. If our queue is empty, every rank's is — no + # indices to negotiate, so we can safely skip the collective. + if not self.grammar_queue: + return [] + + now = time.monotonic() + ready_idxs: set[int] = set() + failed_idxs: set[int] = set() + + for i, (_, state, _) in enumerate(self.grammar_queue): + if state.finished or state.grammar is None: + # Aborted while queued. + ready_idxs.add(i) + continue + + if not isinstance(state.grammar, futures.Future): + raise TypeError( + f"Queued grammar state must hold a Future, got " + f"{type(state.grammar).__name__}: {state=}" + ) + + if state.grammar.done(): + ready_idxs.add(i) + continue + + # Two-phase timeout: + # - while still queued in the executor, bound the queue-wait by + # compile_timeout_secs measured from queued_ts; + # - once the worker actually starts running init_value_impl, + # switch to a fresh compile_timeout_secs budget measured + # from compile_started_at (so executor queueing doesn't eat + # into the per-compile budget). + # Worst-case total wait is ~2 * compile_timeout_secs, but a + # request can never wait forever even if the executor is wedged. + started_at = self.grammar_backend.compile_started_at(state.grammar_key) + + if started_at is None: + elapsed = now - state.grammar_queued_ts + + else: + elapsed = now - started_at + + if elapsed >= self.compile_timeout_secs: + # Closes the race where the worker finishes between the + # state.grammar.done() check above and the elapsed check + # here: prefer admitting a successfully-completed compile + # over aborting it for timeout. + if state.grammar.done(): + ready_idxs.add(i) + + else: + failed_idxs.add(i) + + if self.grammar_sync_size > 1: + gathered: list[tuple[set, set]] = [None] * self.grammar_sync_size + + torch.distributed.all_gather_object( + gathered, + (ready_idxs, failed_idxs), + group=self.grammar_sync_group, + ) + + ready_idxs = set.intersection(*[g[0] for g in gathered]) + failed_idxs = set.union(*[g[1] for g in gathered]) + + if not ready_idxs and not failed_idxs: + return [] + + promoted: list[QueueEntry] = [] + + for i in sorted(ready_idxs): + spec, state, bootstrap = self.grammar_queue[i] + promoted.append((spec, state, bootstrap)) + + if state.finished or state.grammar is None: + continue + + # init_value_impl is supposed to fold compile errors into an + # InvalidGrammarObject, but it can leak (e.g. KeyError from the + # structural_tag legacy dict-walk in xgrammar_backend). Catching + # here keeps the leak from re-raising out of the event loop. + try: + value = state.grammar.result() + except Exception as exc: + value = InvalidGrammarObject(f"{type(exc).__name__}: {exc}") + + if value.is_invalid: + state.grammar = None + + state.set_finish_with_abort( + f"Failed to compile {state.grammar_key[0]} grammar: " + f"{value.error_message}" + ) + + else: + # Future.result() returns the same object to every waiter on a + # shared compile; copy so each request has its own matcher + # state. + state.grammar = value.copy() + + # Dedupe record_compile_timeout by key: multiple concurrent + # requests share a single compile, so one timeout wave would + # otherwise increment the per-key retry counter N times and + # blow past compile_max_retries in a single pass. + timed_out_keys: set = set() + + for i in sorted(failed_idxs): + spec, state, bootstrap = self.grammar_queue[i] + promoted.append((spec, state, bootstrap)) + + # Race: the compile future may have completed successfully + # between the classification loop and this block. Re-check + # done() so we don't falsely time out a compile that just + # finished — the cost of one extra ``result()`` here is + # negligible compared to the hours a user could spend + # debugging a phantom timeout. + if isinstance(state.grammar, futures.Future) and state.grammar.done(): + try: + value = state.grammar.result() + except BaseException: + value = None + + if value is not None and not value.is_invalid: + # See above: copy per-request so concurrent waiters on a + # shared compile future don't share matcher state. + state.grammar = value.copy() + continue + + # Don't cancel the future: it's shared across all + # concurrent requests for this key. The compile runs to + # completion in the background; its result (or the timeout + # marker we post below) lands in the cache for future + # requests to consume. + state.grammar = None + + timeout_msg = ( + f"Grammar compilation timed out after {self.compile_timeout_secs:.1f}s" + ) + + # Cache the timeout so concurrent retries short-circuit + # instead of all timing out one by one. The backend tracks + # attempt count per key: each timeout caches a TTL'd marker, + # and after compile_max_retries timeouts the marker is + # escalated to permanent (the compiler is consistently + # broken for this key). + if state.grammar_key not in timed_out_keys: + timed_out_keys.add(state.grammar_key) + self.grammar_backend.record_compile_timeout( + state.grammar_key, + timeout_msg, + ttl_secs=self.compile_timeout_secs, + max_retries=self.compile_max_retries, + ) + + state.set_finish_with_abort(f"{timeout_msg}: key={state.grammar_key}") + + drop = ready_idxs | failed_idxs + + self.grammar_queue = [ + entry for i, entry in enumerate(self.grammar_queue) if i not in drop + ] + + return promoted diff --git a/python/tokenspeed/runtime/grammar/reasoning_structural_tag.py b/python/tokenspeed/runtime/grammar/reasoning_structural_tag.py new file mode 100644 index 0000000..f17c491 --- /dev/null +++ b/python/tokenspeed/runtime/grammar/reasoning_structural_tag.py @@ -0,0 +1,188 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Compose xgrammar structural tags that defer constraint enforcement +past the model's reasoning channel. + +Given a tokenspeed ``--reasoning-parser`` name and any xgrammar content +format (``JSONSchemaFormat``, ``RegexFormat``, ``GrammarFormat``, …), +``structural_tag_for_reasoning_response()`` produces an xgrammar +structural-tag JSON string that: + +1. Uses xgrammar's ``get_builtin_structural_tag`` template for the + matching model family (so the channel layout matches the model's + chat-template expectations exactly). +2. Substitutes the response-channel content (default ``any_text``) with + the supplied content format, so the constraint only kicks in inside + that channel — reasoning content stays free-form. + +This avoids the channel-preamble corruption that activating any grammar +at token 0 would cause for reasoning models (the model emits +``<|channel|>analysis<|message|>...`` etc. before the response, and a +naïve grammar would reject those tokens). + +Convenience wrappers per content type are provided +(``..._json_schema``, ``..._regex``, ``..._grammar``); call sites that +already have an xgrammar Format instance can call the base function. +""" + +from __future__ import annotations + +from typing import Any + +# smg reasoning_parser name → xgrammar builtin model template. Keys +# must match ``smg::get_available_reasoning_parsers()`` exactly. Unmapped +# parsers fall through to the unwrapped json_schema constraint. +# gpt-oss / harmony is intentionally absent: smg's Harmony Responses path +# wraps the schema in a structural tag itself, so the engine never sees +# json_schema for that family. +_REASONING_PARSER_TO_XGRAMMAR_MODEL: dict[str, str] = { + "deepseek_r1": "deepseek_r1", + "deepseek_v31": "deepseek_v3_2", + "kimi": "kimi", + "kimi_k25": "kimi", + "minimax": "minimax", + "qwen3": "qwen_3", + "qwen3_thinking": "qwen_3", +} + + +def structural_tag_for_reasoning_response( + reasoning_parser: str, content_format: Any +) -> str | None: + """Build a JSON-serialized xgrammar structural tag for ``reasoning + content``. + + ``content_format`` is any xgrammar content Format instance — + ``JSONSchemaFormat``, ``RegexFormat``, ``GrammarFormat``, etc. It + replaces the response channel's default ``any_text``. + + Returns ``None`` if ``reasoning_parser`` has no xgrammar mapping — + the caller should leave the user's constraint unchanged in that case. + """ + model = _REASONING_PARSER_TO_XGRAMMAR_MODEL.get(reasoning_parser) + if model is None: + return None + + import xgrammar + + tag = xgrammar.get_builtin_structural_tag(model=model, reasoning=True) + fmt = tag.format + + # Track whether the substitution actually landed. If it didn't + # (template shape changed in a future xgrammar release, response + # slot moved, etc.) we MUST NOT return the tag with the user's + # constraint silently dropped — return None so the caller falls + # back to the unwrapped constraint and the user gets the + # constraint they asked for (at the cost of channel-preamble + # corruption on reasoning models). + substituted = False + + if fmt.type == "tags_with_separator": + # harmony: the builtin template uses ``tags_with_separator`` + # which allows the model to repeat ``analysis sep final`` + # cycles indefinitely — and after the first JSON it does, + # producing trailing ``<|end|><|start|>assistant…`` plus a + # second JSON that the reasoning parser concatenates, + # breaking ``json.loads``. ``stop_after_first=True`` is + # *worse*: it ends after the FIRST tag (analysis only), so + # ``content`` ends up empty. + # + # Instead we replace the format with a fixed sequence of + # exactly one analysis tag, the literal separator, and one + # final tag carrying the user's content. The matcher + # forbids any further channel re-open after the final tag's + # ``<|end|>``, so the model stops cleanly. + from xgrammar.structural_tag import ( + ConstStringFormat, + SequenceFormat, + TagFormat, + ) + + analysis_tag: TagFormat | None = None + final_tag: TagFormat | None = None + for t in fmt.tags: + if t.begin == "<|channel|>analysis<|message|>": + analysis_tag = t + elif t.begin == "<|channel|>final<|message|>": + final_tag = t + if analysis_tag is not None and final_tag is not None: + final_tag.content = content_format + tag.format = SequenceFormat( + elements=[ + analysis_tag, + ConstStringFormat(value=fmt.separator), + final_tag, + ] + ) + substituted = True + elif fmt.type == "sequence": + # qwen / deepseek_r1 / glm47 / kimi / minimax: the layout is + # ``[reasoning_tag, response_any_text]``. Anchor on the LAST + # element (the response) so the reasoning tag is never replaced + # — even if xgrammar later inserts pre-reasoning preamble + # elements. + last_idx = len(fmt.elements) - 1 + if ( + last_idx >= 0 + and getattr(fmt.elements[last_idx], "type", None) == "any_text" + ): + fmt.elements[last_idx] = content_format + substituted = True + + if not substituted: + # Either the format type isn't one we know how to handle, or + # the response slot wasn't where we expected. Bail out so the + # user's constraint isn't silently dropped. + return None + + return tag.model_dump_json(by_alias=True) + + +def structural_tag_for_reasoning_json_schema( + reasoning_parser: str, user_schema: Any +) -> str | None: + """Convenience wrapper: response channel constrained to a JSON schema.""" + from xgrammar.structural_tag import JSONSchemaFormat + + return structural_tag_for_reasoning_response( + reasoning_parser, JSONSchemaFormat(json_schema=user_schema) + ) + + +def structural_tag_for_reasoning_regex( + reasoning_parser: str, pattern: str +) -> str | None: + """Convenience wrapper: response channel constrained to a regex pattern.""" + from xgrammar.structural_tag import RegexFormat + + return structural_tag_for_reasoning_response( + reasoning_parser, RegexFormat(pattern=pattern) + ) + + +def structural_tag_for_reasoning_grammar( + reasoning_parser: str, ebnf_grammar: str +) -> str | None: + """Convenience wrapper: response channel constrained to an EBNF grammar.""" + from xgrammar.structural_tag import GrammarFormat + + return structural_tag_for_reasoning_response( + reasoning_parser, GrammarFormat(grammar=ebnf_grammar) + ) diff --git a/python/tokenspeed/runtime/grammar/xgrammar_backend.py b/python/tokenspeed/runtime/grammar/xgrammar_backend.py new file mode 100755 index 0000000..c0ae5c7 --- /dev/null +++ b/python/tokenspeed/runtime/grammar/xgrammar_backend.py @@ -0,0 +1,214 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright 2023-2024 SGLang Team +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Constrained decoding with xgrammar backend.""" + +import json + +import torch +from xgrammar import ( + CompiledGrammar, + GrammarCompiler, + GrammarMatcher, + StructuralTagItem, + TokenizerInfo, + allocate_token_bitmask, +) + +from tokenspeed.runtime.grammar.base_grammar_backend import ( + BaseGrammarBackend, + BaseGrammarObject, + InvalidGrammarObject, +) +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +MAX_ROLLBACK_TOKENS = 200 + + +class XGrammarGrammar(BaseGrammarObject): + def __init__( + self, + matcher: GrammarMatcher, + vocab_size: int, + ctx: CompiledGrammar, + override_stop_tokens: list[int] | int | None, + ) -> 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: list[int] = [] + + def is_terminated(self): + + return self.matcher.is_terminated() + + def accept_token(self, token: int): + + if not self.is_terminated(): + if not self.matcher.accept_token(token): + raise ValueError( + f"Tokens not accepted: {token}\n" + f"Accepted tokens: {self.accepted_tokens}\n" + f"Terminated: {self.matcher.is_terminated()}\n" + ) + + else: + self.accepted_tokens.append(token) + + def try_accept_token(self, token: int) -> bool: + """Non-raising accept used by the spec-verify mask fill. + + Returns True iff the matcher accepts the token; on rejection the + matcher state is unchanged. Used to walk draft chains where some + positions may diverge from the grammar. + """ + if self.is_terminated(): + return False + + else: + return self.matcher.accept_token(token) + + def rollback(self, k: int): + + self.matcher.rollback(k) + self.accepted_tokens = self.accepted_tokens[:-k] + + 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 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 + ) + + +class XGrammarGrammarBackend(BaseGrammarBackend): + def __init__( + self, + tokenizer, + vocab_size: int, + disable_any_whitespace: bool = False, + ) -> None: + + super().__init__() + + tokenizer_info = TokenizerInfo.from_huggingface( + tokenizer, vocab_size=vocab_size + ) + + self.grammar_compiler = GrammarCompiler(tokenizer_info=tokenizer_info) + + self.vocab_size = vocab_size + self.override_stop_tokens = None + self.disable_any_whitespace = disable_any_whitespace + + def init_value_impl( + self, key: tuple[str, str], require_reasoning: bool + ) -> BaseGrammarObject: + + key_type, key_string = key + any_whitespace = not self.disable_any_whitespace + try: + if key_type == "json": + if key_string == "$$ANY$$": + ctx = self.grammar_compiler.compile_json_schema( + '{"type": "object"}', any_whitespace=any_whitespace + ) + else: + ctx = self.grammar_compiler.compile_json_schema( + schema=key_string, any_whitespace=any_whitespace + ) + + elif key_type == "ebnf": + ctx = self.grammar_compiler.compile_grammar(key_string) + + elif key_type == "regex": + ctx = self.grammar_compiler.compile_regex(key_string) + + elif key_type == "structural_tag": + structural_tag = json.loads(key_string) + + # Built-in structural-tag payloads include a ``format`` field + # and can be compiled directly. Explicit structures/triggers + # payloads are expanded into xgrammar tag items below. + if "format" in structural_tag: + ctx = self.grammar_compiler.compile_structural_tag(structural_tag) + else: + 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"] + ) + + else: + raise ValueError(f"Invalid key_type: {key_type}") + + except (RuntimeError, ValueError, json.JSONDecodeError) as exc: + logger.warning( + "Failed to compile %s grammar: key_string=%r, e=%r", + key_type, + key_string, + exc, + ) + return InvalidGrammarObject(f"{type(exc).__name__}: {exc}") + + matcher = GrammarMatcher(ctx, max_rollback_tokens=MAX_ROLLBACK_TOKENS) + return XGrammarGrammar(matcher, self.vocab_size, ctx, self.override_stop_tokens) + + def reset(self): + + self.grammar_compiler and self.grammar_compiler.clear_cache() diff --git a/python/tokenspeed/runtime/layers/activation.py b/python/tokenspeed/runtime/layers/activation.py new file mode 100755 index 0000000..a3b271b --- /dev/null +++ b/python/tokenspeed/runtime/layers/activation.py @@ -0,0 +1,124 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Fused operators for activation layers.""" + +from dataclasses import dataclass + +import torch +import triton +import triton.language as tl +from tokenspeed_kernel.platform import current_platform + +from tokenspeed.runtime.utils import ( + get_colorful_logger, +) +from tokenspeed.runtime.utils.pdl import pdl_enabled + +_is_amd = current_platform().is_amd + +if _is_amd: + from tokenspeed_kernel.ops.activation.triton import silu_and_mul + +else: + from tokenspeed_kernel.ops.activation.flashinfer import ( + silu_and_mul, + ) + +logger = get_colorful_logger(__name__) + + +class SiluAndMul(torch.nn.Module): + + def forward(self, x: torch.Tensor, fp8_out: bool = False) -> torch.Tensor: + if not _is_amd: + + def get_tma_aligned_scale(x): + aligned_size = (x.shape[-2] + 3) // 4 * 4 + x_s = torch.empty( + x.shape[:-2] + (x.shape[-1] // 128, aligned_size), + device=x.device, + dtype=torch.float32, + ).permute(-1, -2)[: x.shape[-2], :] + return x_s + + d = x.shape[-1] // 2 + output_shape = x.shape[:-1] + (d,) + if fp8_out: + out = torch.empty( + output_shape, dtype=torch.float8_e4m3fn, device=x.device + ) + scale = get_tma_aligned_scale(out) + from tokenspeed_kernel.ops.activation.cuda import ( + silu_and_mul_fuse_block_quant, + ) + + out, scale = silu_and_mul_fuse_block_quant( + x, scale, out, enable_pdl=pdl_enabled() + ) + return out, scale + else: + out = torch.empty(output_shape, dtype=x.dtype, device=x.device) + silu_and_mul(x, out, enable_pdl=pdl_enabled()) + return out + + if fp8_out: + raise NotImplementedError("AMD fp8_out silu_and_mul is not implemented") + d = x.shape[-1] // 2 + out = torch.empty(x.shape[:-1] + (d,), dtype=x.dtype, device=x.device) + return silu_and_mul(x, out, enable_pdl=pdl_enabled()) + + +@triton.jit +def clip(x, limit, clip_lower: tl.constexpr): + res = tl.minimum(x, limit) + if clip_lower: + res = tl.maximum(-limit, res) + return res + + +@triton.jit +def compute_swiglu(gelu, linear, scale, alpha, limit): + gelu = gelu.to(tl.float32) * scale + if limit is not None: + gelu = clip(gelu, limit, clip_lower=False) + linear = linear.to(tl.float32) * scale + if limit is not None: + linear = clip(linear, limit, clip_lower=True) + + s = gelu / (1 + tl.exp(-alpha * gelu)) + + return tl.fma(s, linear, s) # (s * (linear + 1)) + + +@triton.jit(repr=lambda _: "_swiglu") +def swiglu_fn(input, alpha, limit, exclusive_sum, local_num_experts): + begin = exclusive_sum[0] + end = exclusive_sum[local_num_experts] + input = input[begin:end] + + gelu, linear = tl.split(tl.reshape(input, (input.shape[0], input.shape[1] // 2, 2))) + return compute_swiglu(gelu, linear, 1.0, alpha, limit) + + +@dataclass +class SwigluArg: + alpha: float + limit: float diff --git a/python/tokenspeed/runtime/layers/attention/backends/__init__.py b/python/tokenspeed/runtime/layers/attention/backends/__init__.py new file mode 100644 index 0000000..6914165 --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/backends/__init__.py @@ -0,0 +1,42 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +# ruff: noqa: E402,F401 +# Import all backend modules to trigger register_backend() calls. +import logging + +from tokenspeed_kernel.platform import current_platform + +platform = current_platform() +logger = logging.getLogger(__name__) + + +if platform.is_nvidia: + from tokenspeed.runtime.layers.attention.backends import deepseek_v4 # noqa: F401 + from tokenspeed.runtime.layers.attention.backends import flashmla # noqa: F401 + from tokenspeed.runtime.layers.attention.backends import trtllm # noqa: F401 + from tokenspeed.runtime.layers.attention.backends import trtllm_mla # noqa: F401 + from tokenspeed.runtime.layers.attention.backends import ( # noqa: F401 + tokenspeed_mla, + ) + +from tokenspeed.runtime.layers.attention.backends import dsa # noqa: F401 +from tokenspeed.runtime.layers.attention.backends import mha # noqa: F401 +from tokenspeed.runtime.layers.attention.backends import mla # noqa: F401 diff --git a/python/tokenspeed/runtime/layers/attention/backends/base.py b/python/tokenspeed/runtime/layers/attention/backends/base.py new file mode 100644 index 0000000..cfd618b --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/backends/base.py @@ -0,0 +1,291 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import inspect +from abc import ABC, abstractmethod +from contextlib import contextmanager +from typing import TYPE_CHECKING + +import torch + +from tokenspeed.runtime.execution.breakable_cuda_graph import break_point + +if TYPE_CHECKING: + from tokenspeed.runtime.execution.forward_batch_info import ForwardMode + from tokenspeed.runtime.layers.attention.configs.base import BaseAttnConfig + from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool + from tokenspeed.runtime.layers.paged_attention import PagedAttention + from tokenspeed.runtime.pd.utils import StepCounter + + +def init_backend_cuda_graph_state( + backend: "AttentionBackend", + max_bs: int, + seq_lens_buf: torch.Tensor, + **extras, +) -> None: + """Call ``backend.init_cuda_graph_state`` with only the kwargs its + signature accepts (VAR_KEYWORD accepts all of them). + + Signature-probe instead of try/except TypeError: paged_cache_group_specs + is load-bearing for the state shed, so a TypeError raised from inside the + backend's body must propagate rather than silently retry without specs. + + Shared by the cuda-graph wrapper and by composite backends (hybrid) that + forward to user-selectable sub-backends with possibly narrow signatures. + """ + params = inspect.signature(backend.init_cuda_graph_state).parameters + if not any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()): + extras = {k: v for k, v in extras.items() if k in params} + backend.init_cuda_graph_state(max_bs, seq_lens_buf, **extras) + + +class AttentionBackend(ABC): + """The base class of attention backends""" + + uses_paged_cache_groups: bool = False + # Flat KV-cache per-group block tables (absolute index, null hole = 0). A + # separate flag from uses_paged_cache_groups because the two mechanisms have + # different hole/index semantics; a group-aware flat backend (Phase 4) sets + # this True. Default False keeps every existing backend on today's path. + uses_flat_cache_groups: bool = False + # False for flat-capable backends whose spec-verify path is not wired yet. + flat_spec_capable: bool = True + uses_padded_decode_token_mask: bool = False + + def __init__(self, config: BaseAttnConfig) -> None: + self.device = config.device + self.num_qo_heads = config.num_attention_heads // config.attn_tp_size + self.num_kv_heads = max(config.num_kv_heads // config.attn_tp_size, 1) + self.dtype = config.dtype + self.head_dim = config.head_dim + self.is_draft = config.is_draft + self.spec_num_tokens = config.speculative_num_draft_tokens + # True when this backend's CUDA-graph block-table (kv_indices) buffer is + # aliased to a peer backend's (e.g. a drafter sharing the target's), so + # the replay path skips rebuilding it — the peer already populates it. + self._block_table_aliased = False + + @contextmanager + def override_num_extends(self, num_extends: int): + """Temporarily override the decode-metadata slice discriminator for the + wrapped block. Used by MLA backends to flip between drafter step 0 + (slice = [num_extends:]) and step 1+ (slice = [0:]). + + Default no-op for backends that fill separate prefill/decode metadata + at init time. + """ + yield + + def support_kv_cache_prewrite( + self, forward_mode: ForwardMode | None = None + ) -> bool: + return False + + def select_out_cache_loc(self, layer, out_cache_loc, forward_mode=None): + """Flat per-group write-location hook for out-of-backend KV writers + (fused RoPE prewrite); identity for backends without flat cache + groups (see uses_flat_cache_groups). ``forward_mode`` picks the + metadata slot for backends that prewrite on extend as well.""" + return out_cache_loc + + @property + def sinks_dtype(self) -> torch.dtype: + return torch.bfloat16 + + @abstractmethod + def init_forward_metadata(self, *args, **kwargs): + """Init the metadata for a forward pass. + + When use_cuda_graph=True the backend should use its pre-allocated + cuda-graph buffers instead of the normal eager buffers. + """ + raise NotImplementedError() + + def init_cuda_graph_state(self, max_bs: int, seq_lens_buf: torch.Tensor): + """Init the global shared states for cuda graph. `seq_lens_buf` is + the controller-owned per-request seq_lens; backends should reference + (alias) it rather than copy, and must not mutate the contents.""" + raise NotImplementedError() + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + forward_mode: ForwardMode, + flat_cache_group_ids: tuple[str, ...] = (), + **kwargs, + ): + """Init the metadata for a forward pass for capturing a cuda graph. + + ``flat_cache_group_ids`` names the flat KV-cache groups whose page + tables arrive at replay; a flat-capable backend (uses_flat_cache_groups) + allocates its persistent per-group buffers from these ids — no table + data exists at capture time. Empty tuple for non-flat backends. + """ + raise NotImplementedError() + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + forward_mode: ForwardMode = None, + req_to_page: torch.Tensor = None, + flat_block_tables: dict[str, torch.Tensor] | None = None, + **kwargs, + ): + """Update pre-allocated CUDA-graph metadata buffers in-place before replay. + + Called instead of init_forward_metadata when use_cuda_graph=True, so + that the captured kernels (which hold pointers into the pre-allocated + buffers) see the current batch's data without any new allocations. + ``flat_block_tables`` carries the per-group flat page tables + (group_id -> [>=bs, cols]) for flat-capable backends; a backend that + captured flat buffers must be handed non-empty tables whenever bs > 0. + Default: fall back to init_forward_metadata (correct but may not work + for all backends that use separate cuda-graph buffer pools). + """ + raise NotImplementedError( + f"{type(self).__name__} must implement init_forward_metadata_replay_cuda_graph " + "for CUDA graph support" + ) + + def configure_runtime(self, **kwargs) -> None: + """Configure runtime state after model loading (e.g. sliding_window_size). + + Called once during ModelExecutor initialization with information that is + not available at backend construction time. Default: no-op. + """ + pass + + def register_step_counter(self, step_counter: StepCounter): + self.step_counter = step_counter + + @contextmanager + def record_pd_cache_step( + self, + forward_mode: ForwardMode, + save_kv_cache: bool, + record_kv_cache: bool | None, + ): + """Anchor the PD layerwise cache-step record to the wrapped KV write. + + Records the ``StepCounter`` step before the attention call when the KV + was pre-written (``save_kv_cache=False``) and after it otherwise, so a + layerwise cache transfer always observes a fully written layer. See + ``forward`` for the ``record_kv_cache`` override contract. No-op when no + step counter is registered. Backends that own the record (e.g. the + hybrid wrapper, which counts once per model layer across full-attn + + mamba children) reuse this to avoid duplicating the gate logic. + """ + if record_kv_cache is None: + record_cache = not forward_mode.is_decode() and not forward_mode.is_idle() + else: + record_cache = record_kv_cache + record_cache = record_cache and getattr(self, "step_counter", None) is not None + + if record_cache and not save_kv_cache: + self.step_counter.record_cache() + yield + if record_cache and save_kv_cache: + self.step_counter.record_cache() + + @break_point + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: PagedAttention, + out_cache_loc: torch.Tensor, + token_to_kv_pool: BaseTokenToKVPool, + forward_mode: ForwardMode, + bs: int, + save_kv_cache: bool = True, + record_kv_cache: bool | None = None, + **kwargs, + ): + """Run forward on an attention layer with explicit scheduler metadata. + + ``record_kv_cache`` overrides the PD layerwise cache-step recording: + ``None`` keeps the default (record on the EXTEND-side path), an explicit + bool forces it so a DECODE-dispatched draft catch-up can still record. + """ + with self.record_pd_cache_step(forward_mode, save_kv_cache, record_kv_cache): + if forward_mode.is_decode(): + ret = self.forward_decode( + q, + k, + v, + layer, + out_cache_loc, + token_to_kv_pool, + bs, + save_kv_cache=save_kv_cache, + **kwargs, + ) + else: + ret = self.forward_extend( + q, + k, + v, + layer, + out_cache_loc, + token_to_kv_pool, + bs, + save_kv_cache=save_kv_cache, + forward_mode=forward_mode, + **kwargs, + ) + return ret + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: PagedAttention, + out_cache_loc: torch.Tensor, + token_to_kv_pool: BaseTokenToKVPool, + bs: int, + save_kv_cache: bool = True, + **kwargs, + ): + """Run a forward for decode.""" + raise NotImplementedError() + + def forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: PagedAttention, + out_cache_loc: torch.Tensor, + token_to_kv_pool: BaseTokenToKVPool, + bs: int, + save_kv_cache: bool = True, + **kwargs, + ): + """Run a forward for extend.""" + raise NotImplementedError() diff --git a/python/tokenspeed/runtime/layers/attention/backends/deepseek_v4.py b/python/tokenspeed/runtime/layers/attention/backends/deepseek_v4.py new file mode 100644 index 0000000..a8acb79 --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/backends/deepseek_v4.py @@ -0,0 +1,2098 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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 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. + +from __future__ import annotations + +import torch +from tokenspeed_kernel.ops.attention.flash_mla import ( + flash_mla_sparse_fwd, + flash_mla_with_kvcache, + get_mla_metadata, +) +from tokenspeed_kernel.ops.attention.triton.deepseek_v4 import ( + deepseek_v4_indexer_decode_metadata_compute, +) +from tokenspeed_kernel.registry import error_fn + +try: + from tokenspeed_kernel.thirdparty import deep_gemm +except Exception: + deep_gemm = None # type: ignore[assignment] + +from tokenspeed.runtime.configs.deepseek_v4_cache_spec import ( + DEEPSEEK_V4_SPARSE_PREFILL_TOPK_ALIGNMENT, + deepseek_v4_swa_row_bytes, + v4_compressed_kv_group_id, +) +from tokenspeed.runtime.configs.model_config import AttentionArch +from tokenspeed.runtime.configs.paged_cache_spec import ( + compute_max_logical_pages_for_capture, +) +from tokenspeed.runtime.execution.forward_batch_info import ForwardMode +from tokenspeed.runtime.layers.attention.backends.base import AttentionBackend +from tokenspeed.runtime.layers.attention.deepseek_v4.metadata import ( + DeepseekV4ForwardMetadata, +) +from tokenspeed.runtime.layers.attention.deepseek_v4_ops import ( + deepseek_v4_build_dense_prefill_local_compressed_indices, + deepseek_v4_combine_dense_swa_indices, + deepseek_v4_combine_topk_swa_indices, + deepseek_v4_compute_global_topk_indices_and_lens, + deepseek_v4_decode_swa_indices_and_lens, + deepseek_v4_dequantize_and_gather_k_cache, +) +from tokenspeed.runtime.layers.attention.kv_cache.deepseek_v4 import ( + DeepseekV4CacheMetadata, + _split_paged_cache_block_tables_into_v4_metadata, +) +from tokenspeed.runtime.layers.attention.registry import register_backend +from tokenspeed.runtime.utils.env import global_server_args_dict +from tokenspeed.runtime.utils.nvtx import nvtx_range + +DEEPSEEK_V4_DEFAULT_PREFILL_CHUNK_SIZE = 4 + + +def _compressed_block_table_base_offsets( + metadata: DeepseekV4ForwardMetadata, + compress_ratio: int, +) -> torch.Tensor | None: + return metadata.cache.paged_cache_block_table_base_offsets.get( + v4_compressed_kv_group_id(compress_ratio) + ) + + +def _decode_positions_from_metadata( + metadata: DeepseekV4ForwardMetadata, + num_tokens: int, +) -> torch.Tensor: + token_to_req = metadata.token_to_req_indices[:num_tokens].to(torch.int64) + query_starts = metadata.query_start_loc[token_to_req].to(torch.int64) + query_lens = metadata.query_lens[token_to_req].to(torch.int64) + seq_lens = metadata.seq_lens[token_to_req].to(torch.int64) + token_offsets = torch.arange( + num_tokens, + dtype=torch.int64, + device=metadata.seq_lens.device, + ) + return seq_lens - query_lens + token_offsets - query_starts + + +def _refresh_decode_indexer_plan_cache( + metadata: DeepseekV4ForwardMetadata, + *, + max_context_len: int, +) -> None: + """Pre-build decode-indexer plan tensors before per-layer parallel work. + + This keeps per-layer indexer calls read-only with respect to cached plan + buffers while compressor work may run on an auxiliary stream. + """ + indexer_metadata = metadata.indexer + cache = indexer_metadata.decode_plan_cache + if not cache: + return + refreshed_keys = indexer_metadata.decode_plan_refreshed_keys + refreshed_keys.clear() + for ( + compress_ratio, + cache_block_size, + num_tokens, + ), plan in list(cache.items()): + if num_tokens <= 0: + plan.context_lens.zero_() + plan.block_table.zero_() + plan.max_context_len = 0 + refreshed_keys.add((compress_ratio, cache_block_size, num_tokens)) + continue + positions = _decode_positions_from_metadata(metadata, num_tokens) + token_to_req_indices = metadata.token_to_req_indices[:num_tokens] + block_table = metadata.cache.compressed_block_table( + compress_ratio, + cache_block_size, + ) + block_table_base_offsets = ( + _compressed_block_table_base_offsets(metadata, compress_ratio) + if block_table is not metadata.cache.block_table + else None + ) + rows = int(block_table.shape[0]) if block_table.ndim >= 1 else 0 + cols = int(block_table.shape[1]) if block_table.ndim >= 2 else 0 + if rows <= 0 or cols <= 0: + plan.context_lens.zero_() + plan.block_table.zero_() + plan.max_context_len = 0 + refreshed_keys.add((compress_ratio, cache_block_size, num_tokens)) + continue + max_blocks = int(plan.block_table.shape[1]) + if max_context_len > 0: + derived_max_len = max( + 1, + (max_context_len + compress_ratio - 1) // compress_ratio, + ) + else: + derived_max_len = max( + 1, + (block_table.shape[1] * cache_block_size + compress_ratio - 1) + // compress_ratio, + ) + if plan.max_context_len != derived_max_len: + plan.max_context_len = derived_max_len + deepseek_v4_indexer_decode_metadata_compute( + positions=positions, + token_to_req_indices=token_to_req_indices, + block_table=block_table, + cache_block_size=cache_block_size, + compress_ratio=compress_ratio, + max_blocks=max_blocks, + out_context_lens=plan.context_lens, + out_block_tables=plan.block_table, + block_table_base_offsets=block_table_base_offsets, + ) + if metadata.is_valid_token is not None: + valid = metadata.is_valid_token[:num_tokens].to( + device=plan.context_lens.device, + dtype=torch.bool, + ) + with torch.inference_mode(): + plan.context_lens.masked_fill_(~valid.view(num_tokens, 1), 0) + plan.block_table.masked_fill_( + ~valid.to(device=plan.block_table.device).view(num_tokens, 1), + 0, + ) + refreshed_keys.add((compress_ratio, cache_block_size, num_tokens)) + + +def _refresh_decode_indexer_schedule_metadata( + metadata: DeepseekV4ForwardMetadata, +) -> None: + indexer_metadata = metadata.indexer + if not indexer_metadata.decode_schedule_metadata_cache: + return + if deep_gemm is None: + return + get_metadata = getattr(deep_gemm, "get_paged_mqa_logits_metadata", None) + if get_metadata is None: + return + for ( + compress_ratio, + cache_block_size, + num_tokens, + ), schedule_metadata in list( + indexer_metadata.decode_schedule_metadata_cache.items() + ): + if num_tokens <= 0: + continue + key = (compress_ratio, cache_block_size, num_tokens) + decode_plan = indexer_metadata.decode_plan_cache.get(key) + context_lens = getattr(decode_plan, "context_lens", None) + if ( + context_lens is not None + and context_lens.shape == (num_tokens, 1) + and context_lens.dtype == torch.int32 + ): + context_lens = context_lens.contiguous() + else: + positions = _decode_positions_from_metadata(metadata, num_tokens) + compressed_lens = torch.div( + positions.to(torch.int32) + 1, + compress_ratio, + rounding_mode="floor", + ).clamp_min(0) + if metadata.is_valid_token is not None: + valid = metadata.is_valid_token[:num_tokens].to( + device=compressed_lens.device, + dtype=torch.bool, + ) + compressed_lens = torch.where( + valid, + compressed_lens, + torch.zeros_like(compressed_lens), + ) + context_lens = compressed_lens.view(num_tokens, 1).contiguous() + refreshed = get_metadata( + context_lens, + cache_block_size, + deep_gemm.get_num_sms(), + ) + if ( + schedule_metadata.shape == refreshed.shape + and schedule_metadata.device == refreshed.device + and schedule_metadata.dtype == refreshed.dtype + ): + with torch.inference_mode(): + schedule_metadata.copy_(refreshed) + else: + indexer_metadata.decode_schedule_metadata_cache[key] = refreshed + + +class DeepseekV4AttentionBackend(AttentionBackend): + """Metadata owner for the model-local DeepSeek V4 attention path.""" + + uses_paged_cache_groups = True + uses_padded_decode_token_mask = True + + def __init__(self, config) -> None: + super().__init__(config) + self.page_size = config.page_size + self.context_len = config.context_len + rope_head_dim = getattr(config, "qk_rope_head_dim", None) + self._fp8_ds_mla_row_bytes = ( + deepseek_v4_swa_row_bytes(config.head_dim, rope_head_dim) + if rope_head_dim is not None + else None + ) + prefill_chunk_size = getattr(config, "deepseek_v4_prefill_chunk_size", None) + if prefill_chunk_size is None: + prefill_chunk_size = global_server_args_dict.get( + "deepseek_v4_prefill_chunk_size", + DEEPSEEK_V4_DEFAULT_PREFILL_CHUNK_SIZE, + ) + self.prefill_chunk_size = max(1, int(prefill_chunk_size)) + self.max_num_pages = max( + 1, + (self.context_len + self.page_size - 1) // self.page_size, + ) + self.forward_metadata: DeepseekV4ForwardMetadata | None = None + self.forward_prefill_metadata: DeepseekV4ForwardMetadata | None = None + self.forward_decode_metadata: DeepseekV4ForwardMetadata | None = None + self._decode_tile_metadata = {} + self._cuda_graph_metadata = {} + self._cuda_graph_paged_cache_block_tables: dict[str, torch.Tensor] = {} + # Per-sliding-group [max_bs] int32 buffers mirroring the block-table + # buffers; populated by init_cuda_graph_state. + self._cuda_graph_paged_cache_base_offsets: dict[str, torch.Tensor] = {} + self._cuda_graph_max_bs = 0 + self._prefill_workspace_buffer: torch.Tensor | None = None + self._prefill_workspace_rows = 0 + self._prefill_workspace_head_dim = 0 + self._prefill_dense_compressed_indices_buffer: torch.Tensor | None = None + self._decode_swa_window_size = 0 + self._decode_swa_block_size = 0 + self.speculative_num_steps = getattr(config, "speculative_num_steps", 0) or 0 + self.speculative_num_draft_tokens = ( + getattr(config, "speculative_num_draft_tokens", 0) or 0 + ) + self._draft_decode_step = 0 + self._draft_decode_base_seq_lens: torch.Tensor | None = None + self._draft_decode_metadata: DeepseekV4ForwardMetadata | None = None + self._cuda_graph_draft_decode_metadata = {} + self._cuda_graph_query_start_by_tokens_per_req: dict[int, torch.Tensor] = {} + self._cuda_graph_token_to_req_by_tokens_per_req: dict[int, torch.Tensor] = {} + + def _get_prefill_workspace( + self, + *, + num_reqs: int, + workspace_width: int, + head_dim: int, + device: torch.device, + ) -> torch.Tensor: + workspace_reqs = max(1, num_reqs) + rows = workspace_reqs * workspace_width + needs_alloc = ( + self._prefill_workspace_buffer is None + or self._prefill_workspace_buffer.device != device + or self._prefill_workspace_head_dim != head_dim + or self._prefill_workspace_rows < rows + ) + if needs_alloc: + self._prefill_workspace_buffer = torch.empty( + (rows, head_dim), + dtype=torch.bfloat16, + device=device, + ) + self._prefill_workspace_rows = rows + self._prefill_workspace_head_dim = head_dim + assert self._prefill_workspace_buffer is not None + return self._prefill_workspace_buffer[:rows].view( + workspace_reqs, + workspace_width, + head_dim, + ) + + def _query_lens( + self, + bs: int, + num_tokens: int, + seq_lens: torch.Tensor, + forward_mode: ForwardMode | None, + num_extends: int, + extend_seq_lens_cpu: torch.Tensor | None, + extend_prefix_lens_cpu: torch.Tensor | None, + extend_prefix_lens: torch.Tensor | None, + ) -> torch.Tensor: + if forward_mode is not None and forward_mode.is_decode_or_idle(): + if forward_mode.is_decode() and num_tokens != bs: + if bs == 0: + return torch.zeros(0, dtype=torch.int32, device=seq_lens.device) + if num_tokens % bs != 0: + raise RuntimeError( + "DeepSeek V4 packed decode metadata expects uniformly " + f"packed tokens per request, got num_tokens={num_tokens}, " + f"bs={bs}" + ) + tokens_per_req = num_tokens // bs + return torch.full( + (bs,), + tokens_per_req, + dtype=torch.int32, + device=seq_lens.device, + ) + return torch.ones(bs, dtype=torch.int32, device=seq_lens.device) + if forward_mode is not None and forward_mode.is_mixed(): + verify_width = max(1, int(self.speculative_num_draft_tokens)) + lens = torch.full( + (bs,), + verify_width, + dtype=torch.int32, + device=seq_lens.device, + ) + num_prefill_reqs = max(0, min(int(num_extends), bs)) + if num_prefill_reqs == 0: + return lens + if extend_seq_lens_cpu is not None and extend_seq_lens_cpu.numel() > 0: + lens[:num_prefill_reqs] = extend_seq_lens_cpu[:num_prefill_reqs].to( + seq_lens.device, dtype=torch.int32 + ) + elif extend_prefix_lens_cpu is not None: + prefix = extend_prefix_lens_cpu[:num_prefill_reqs].to( + seq_lens.device, dtype=torch.int32 + ) + lens[:num_prefill_reqs] = ( + seq_lens[:num_prefill_reqs].to(torch.int32) - prefix + ).clamp_min(0) + elif extend_prefix_lens is not None: + prefix = extend_prefix_lens[:num_prefill_reqs].to(torch.int32) + lens[:num_prefill_reqs] = ( + seq_lens[:num_prefill_reqs].to(torch.int32) - prefix + ).clamp_min(0) + else: + lens[:num_prefill_reqs] = seq_lens[:num_prefill_reqs].to(torch.int32) + return lens + if extend_seq_lens_cpu is not None: + return extend_seq_lens_cpu[:bs].to(seq_lens.device, dtype=torch.int32) + if extend_prefix_lens_cpu is not None: + prefix = extend_prefix_lens_cpu[:bs].to(seq_lens.device, dtype=torch.int32) + return (seq_lens[:bs].to(torch.int32) - prefix).clamp_min(0) + if extend_prefix_lens is not None: + prefix = extend_prefix_lens[:bs].to(torch.int32) + return (seq_lens[:bs].to(torch.int32) - prefix).clamp_min(0) + return seq_lens[:bs].to(torch.int32) + + def _query_lens_cpu( + self, + bs: int, + forward_mode: ForwardMode | None, + num_extends: int, + extend_seq_lens_cpu: torch.Tensor | None, + extend_prefix_lens_cpu: torch.Tensor | None, + ) -> torch.Tensor | None: + if forward_mode is not None and forward_mode.is_decode_or_idle(): + return torch.ones(bs, dtype=torch.int32) + if forward_mode is not None and forward_mode.is_mixed(): + verify_width = max(1, int(self.speculative_num_draft_tokens)) + lens = torch.full((bs,), verify_width, dtype=torch.int32) + num_prefill_reqs = max(0, min(int(num_extends), bs)) + if num_prefill_reqs == 0: + return lens + if extend_seq_lens_cpu is None: + return None + lens[:num_prefill_reqs] = extend_seq_lens_cpu[:num_prefill_reqs].to( + dtype=torch.int32, device="cpu" + ) + return lens + if extend_seq_lens_cpu is not None: + return extend_seq_lens_cpu[:bs].to(dtype=torch.int32, device="cpu") + if extend_prefix_lens_cpu is not None: + return None + return None + + def _draft_decode_is_valid_token( + self, + prefill_metadata: DeepseekV4ForwardMetadata, + ) -> torch.Tensor | None: + if prefill_metadata.is_valid_token is None: + return None + bs = prefill_metadata.req_pool_indices.numel() + return prefill_metadata.is_valid_token[ + prefill_metadata.query_start_loc[:bs].to(torch.int64) + ] + + def _is_cuda_graph_prefill_metadata( + self, + metadata: DeepseekV4ForwardMetadata, + ) -> bool: + bs = metadata.req_pool_indices.numel() + return self._cuda_graph_metadata.get(bs) is metadata + + def _prepare_draft_decode_metadata( + self, + prefill_metadata: DeepseekV4ForwardMetadata, + base_seq_lens: torch.Tensor, + ) -> None: + self.forward_prefill_metadata = prefill_metadata + self._draft_decode_step = 0 + self._draft_decode_base_seq_lens = base_seq_lens + + bs = prefill_metadata.req_pool_indices.numel() + device = prefill_metadata.req_pool_indices.device + is_cuda_graph_metadata = self._is_cuda_graph_prefill_metadata(prefill_metadata) + metadata = ( + self._cuda_graph_draft_decode_metadata.get(bs) + if is_cuda_graph_metadata + else self._draft_decode_metadata + ) + is_valid_token = self._draft_decode_is_valid_token(prefill_metadata) + if ( + metadata is None + or metadata.req_pool_indices.numel() != bs + or metadata.seq_lens.numel() != bs + or metadata.query_lens.numel() != bs + or metadata.token_to_req_indices.numel() != bs + or metadata.req_pool_indices.device != device + ): + query_lens = torch.ones(bs, dtype=torch.int32, device=device) + token_to_req = torch.arange(bs, dtype=torch.int32, device=device) + decode_seq_lens = torch.empty_like(base_seq_lens) + decode_seq_lens.copy_(base_seq_lens) + decode_is_valid_token = None + if is_valid_token is not None: + decode_is_valid_token = torch.empty_like(is_valid_token) + decode_is_valid_token.copy_(is_valid_token) + metadata = DeepseekV4ForwardMetadata( + req_pool_indices=prefill_metadata.req_pool_indices, + seq_lens=decode_seq_lens, + query_lens=query_lens, + query_start_loc=torch.nn.functional.pad( + torch.cumsum( + query_lens.to(torch.int32), + dim=0, + dtype=torch.int32, + ), + (1, 0), + ), + token_to_req_indices=token_to_req, + cache=prefill_metadata.cache, + is_valid_token=decode_is_valid_token, + forward_mode=ForwardMode.DECODE, + ) + if is_cuda_graph_metadata: + self._cuda_graph_draft_decode_metadata[bs] = metadata + self._draft_decode_metadata = metadata + return + + metadata.req_pool_indices = prefill_metadata.req_pool_indices + metadata.cache = prefill_metadata.cache + metadata.seq_lens.copy_(base_seq_lens) + if is_valid_token is None: + metadata.is_valid_token = None + else: + if ( + metadata.is_valid_token is None + or metadata.is_valid_token.shape != is_valid_token.shape + or metadata.is_valid_token.device != is_valid_token.device + ): + metadata.is_valid_token = torch.empty_like(is_valid_token) + metadata.is_valid_token.copy_(is_valid_token) + metadata.num_prefill_reqs = 0 + metadata.num_prefill_tokens = 0 + metadata.forward_mode = ForwardMode.DECODE + # Reuse path: cached decode-indexer plans still describe the previous + # prefill. Refresh after updating seq_lens so draft step 0 does not + # reuse stale context_lens / block_table tensors. + metadata.cache.refresh_decode_compressed_slot_mappings( + token_to_req_indices=metadata.token_to_req_indices, + query_start_loc=metadata.query_start_loc, + seq_lens=metadata.seq_lens, + is_valid_token=metadata.is_valid_token, + ) + _refresh_decode_indexer_plan_cache( + metadata, + max_context_len=self.context_len, + ) + _refresh_decode_indexer_schedule_metadata(metadata) + self._draft_decode_metadata = metadata + + def _select_decode_metadata( + self, + num_tokens: int, + ) -> DeepseekV4ForwardMetadata | None: + for metadata in (self.forward_metadata, self.forward_decode_metadata): + if ( + metadata is not None + and metadata.forward_mode is not None + and metadata.forward_mode.is_decode() + and metadata.token_to_req_indices.numel() == num_tokens + ): + return metadata + return None + + def init_forward_metadata( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + forward_mode: ForwardMode = None, + req_to_page: torch.Tensor = None, + extend_seq_lens_cpu: torch.Tensor | None = None, + extend_prefix_lens_cpu: torch.Tensor | None = None, + extend_prefix_lens: torch.Tensor | None = None, + **kwargs, + ) -> None: + paged_cache_block_tables = kwargs.pop("paged_cache_block_tables", None) or {} + paged_cache_block_table_base_offsets = ( + kwargs.pop("paged_cache_block_table_base_offsets", None) or {} + ) + num_tokens_arg = kwargs.pop("num_tokens", None) + positions = kwargs.get("positions") + num_extends_arg = kwargs.pop("num_extends", None) + num_extends = bs if num_extends_arg is None else int(num_extends_arg) + if num_tokens_arg is not None: + num_tokens = int(num_tokens_arg) + elif isinstance(positions, torch.Tensor): + num_tokens = int(positions.numel()) + else: + num_tokens = bs + del kwargs + device = seq_lens.device + req_pool_indices = req_pool_indices[:bs] + seq_lens = seq_lens[:bs].to(torch.int32) + query_lens = self._query_lens( + bs, + num_tokens, + seq_lens, + forward_mode, + num_extends, + extend_seq_lens_cpu, + extend_prefix_lens_cpu, + extend_prefix_lens, + ) + is_packed_decode = ( + forward_mode is not None and forward_mode.is_decode() and num_tokens != bs + ) + metadata_forward_mode = forward_mode + if forward_mode is not None and forward_mode.is_mixed(): + num_prefill_reqs = max(0, min(num_extends, bs)) + elif forward_mode is not None and forward_mode.is_extend_or_mixed(): + num_prefill_reqs = bs + else: + num_prefill_reqs = 0 + query_lens_cpu = self._query_lens_cpu( + bs, + forward_mode, + num_extends, + extend_seq_lens_cpu, + extend_prefix_lens_cpu, + ) + seq_lens_cpu = None + if extend_prefix_lens_cpu is not None and query_lens_cpu is not None: + seq_lens_cpu = seq_lens[:bs].to(dtype=torch.int32, device="cpu") + prefix_count = min( + int(extend_prefix_lens_cpu.numel()), + ( + num_prefill_reqs + if forward_mode is not None and forward_mode.is_mixed() + else bs + ), + ) + if prefix_count: + seq_lens_cpu[:prefix_count] = ( + extend_prefix_lens_cpu[:prefix_count].to( + dtype=torch.int32, + device="cpu", + ) + + query_lens_cpu[:prefix_count] + ) + elif extend_seq_lens_cpu is not None and forward_mode is not None: + if forward_mode.is_extend(): + seq_lens_cpu = extend_seq_lens_cpu[:bs].to( + dtype=torch.int32, + device="cpu", + ) + elif forward_mode.is_mixed(): + seq_lens_cpu = seq_lens[:bs].to(dtype=torch.int32, device="cpu") + max_seq_len = int(seq_lens.max().item()) if bs else 0 + if forward_mode is not None and forward_mode.is_extend(): + max_seq_len += max(self.speculative_num_steps - 1, 0) + if is_packed_decode: + max_seq_len += max(int(query_lens.max().item()) - 1, 0) + max_pages = (max_seq_len + self.page_size - 1) // self.page_size + if req_to_page is None: + block_table = torch.zeros( + (bs, max(max_pages, 1)), + dtype=torch.int32, + device=device, + ) + else: + block_table = req_to_page[req_pool_indices, : max(max_pages, 1)] + paged_cache_block_tables = { + str(gid): table[:bs].to(device=device, dtype=torch.int32) + for gid, table in paged_cache_block_tables.items() + } + base_offsets_on_device: dict[str, torch.Tensor] = {} + for gid, off in paged_cache_block_table_base_offsets.items(): + if not isinstance(off, torch.Tensor): + raise TypeError( + "DeepSeek V4 paged_cache_block_table_base_offsets values " + f"must be torch.Tensor, got {type(off).__name__} for " + f"group_id={gid!r}" + ) + base_offsets_on_device[str(gid)] = off[:bs].to( + device=device, dtype=torch.int32 + ) + ( + swa_block_table, + compressor_state_block_tables, + indexer_state_block_table, + swa_base, + compressor_state_base, + indexer_state_base, + ) = _split_paged_cache_block_tables_into_v4_metadata( + paged_cache_block_tables, + base_offsets_on_device, + ) + req_ids = torch.arange(bs, device=device, dtype=torch.int32) + token_to_req = torch.repeat_interleave(req_ids, query_lens.clamp_min(0)) + if ( + forward_mode is not None + and forward_mode.is_mixed() + and num_tokens_arg is not None + ): + # numel() reads tensor shape metadata only. Reducing query_lens and + # calling .item() here would synchronize its CUDA stream on every + # eager mixed batch. + metadata_tokens = token_to_req.numel() + if metadata_tokens != num_tokens: + raise RuntimeError( + "DeepSeek V4 mixed metadata token count mismatch: " + f"query_lens describe {metadata_tokens} tokens, packed input has " + f"{num_tokens}" + ) + num_prefill_tokens = ( + int(query_lens[:num_prefill_reqs].sum().item()) if num_prefill_reqs else 0 + ) + query_start_loc = torch.nn.functional.pad( + torch.cumsum(query_lens.to(torch.int32), dim=0, dtype=torch.int32), + (1, 0), + ) + cache_metadata = DeepseekV4CacheMetadata( + page_size=self.page_size, + block_table=block_table, + paged_cache_block_tables=paged_cache_block_tables, + paged_cache_block_table_base_offsets=base_offsets_on_device, + swa_block_table=swa_block_table, + swa_base_logical_page=swa_base, + compressor_state_block_tables=compressor_state_block_tables, + compressor_state_base_logical_pages=compressor_state_base, + indexer_state_block_table=indexer_state_block_table, + indexer_state_base_logical_page=indexer_state_base, + ) + self.forward_metadata = DeepseekV4ForwardMetadata( + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + query_lens=query_lens, + query_start_loc=query_start_loc, + token_to_req_indices=token_to_req, + cache=cache_metadata, + seq_lens_cpu=seq_lens_cpu, + query_lens_cpu=query_lens_cpu, + num_prefill_reqs=num_prefill_reqs, + num_prefill_tokens=num_prefill_tokens, + forward_mode=metadata_forward_mode, + ) + if is_packed_decode: + self.forward_decode_metadata = self.forward_metadata + if getattr(self, "is_draft", False): + self._prepare_draft_decode_metadata( + self.forward_metadata, + seq_lens.clone(), + ) + elif ( + metadata_forward_mode is not None + and metadata_forward_mode.is_decode_or_idle() + ): + self.forward_decode_metadata = self.forward_metadata + if ( + self.forward_prefill_metadata is not None + and self.forward_prefill_metadata.req_pool_indices.numel() + == seq_lens.numel() + ): + self._prepare_draft_decode_metadata( + self.forward_prefill_metadata, + seq_lens.clone(), + ) + elif forward_mode is not None and forward_mode.is_extend_or_mixed(): + self.forward_prefill_metadata = self.forward_metadata + self._decode_tile_metadata = {} + + def _update_decode_swa_metadata( + self, + metadata: DeepseekV4ForwardMetadata, + *, + window_size: int, + block_size: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + attention_metadata = metadata.attention + num_tokens = metadata.token_to_req_indices.shape[0] + needs_alloc = ( + attention_metadata.decode_swa_indices is None + or attention_metadata.decode_swa_lens is None + or attention_metadata.decode_swa_indices.shape + != ( + num_tokens, + window_size, + ) + or attention_metadata.decode_swa_lens.shape != (num_tokens,) + or attention_metadata.decode_swa_indices.device != metadata.seq_lens.device + ) + if needs_alloc: + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "DeepSeek V4 decode SWA metadata must be allocated before " + "CUDA graph capture" + ) + with torch.inference_mode(False): + attention_metadata.decode_swa_indices = torch.empty( + (num_tokens, window_size), + dtype=torch.int32, + device=metadata.seq_lens.device, + ) + attention_metadata.decode_swa_lens = torch.empty( + (num_tokens,), + dtype=torch.int32, + device=metadata.seq_lens.device, + ) + + cache_metadata = metadata.cache + if cache_metadata.swa_block_table is None: + raise RuntimeError("DeepSeek V4 missing paged-cache block table for SWA KV") + swa_block_table = cache_metadata.swa_block_table + indices, lens = deepseek_v4_decode_swa_indices_and_lens( + query_start_loc=metadata.query_start_loc, + seq_lens=metadata.seq_lens, + token_to_req_indices=metadata.token_to_req_indices, + block_table=swa_block_table, + block_table_base_offsets=cache_metadata.swa_base_logical_page, + window_size=window_size, + block_size=block_size, + is_valid_token=metadata.is_valid_token, + out_indices=attention_metadata.decode_swa_indices, + out_lens=attention_metadata.decode_swa_lens, + ) + attention_metadata.decode_swa_indices = indices + attention_metadata.decode_swa_lens = lens + attention_metadata.decode_swa_window_size = window_size + attention_metadata.decode_swa_block_size = block_size + self._decode_swa_window_size = window_size + self._decode_swa_block_size = block_size + return indices, lens + + def _decode_compressed_attention_indices_and_lens( + self, + positions: torch.Tensor, + *, + compress_ratio: int, + block_size: int, + topk_indices: torch.Tensor | None, + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + if compress_ratio <= 1: + return None, None + metadata = self.forward_metadata + if metadata is None: + raise RuntimeError("DeepSeek V4 decode requires forward metadata") + num_tokens = positions.numel() + req_idx = metadata.token_to_req_indices[:num_tokens].to(torch.int64) + block_table = metadata.cache.compressed_block_table(compress_ratio, block_size) + block_table_base_offsets = ( + _compressed_block_table_base_offsets(metadata, compress_ratio) + if block_table is not metadata.cache.block_table + else None + ) + is_valid_token = ( + metadata.is_valid_token[:num_tokens] + if metadata.is_valid_token is not None + else None + ) + capturing = positions.is_cuda and torch.cuda.is_current_stream_capturing() + if compress_ratio == 4: + if topk_indices is None: + raise RuntimeError("DeepSeek V4 CSA decode requires top-k indices") + topk_local = topk_indices + if block_table_base_offsets is not None: + base_slots = block_table_base_offsets.to( + device=topk_indices.device, + dtype=torch.int64, + )[req_idx] * int(block_size) + topk_i64 = topk_indices.to(torch.int64) + topk_local = torch.where( + topk_i64 >= 0, + topk_i64 - base_slots[:, None], + topk_i64, + ).to(topk_indices.dtype) + indices_2d, lens = deepseek_v4_compute_global_topk_indices_and_lens( + topk_indices=topk_local, + token_to_req_indices=metadata.token_to_req_indices[:num_tokens], + block_table=block_table, + block_size=block_size, + is_valid_token=is_valid_token, + ) + return indices_2d.unsqueeze(1), lens + + cache_key = ( + int(compress_ratio), + int(block_size), + int(num_tokens), + int(positions.data_ptr()) if positions.numel() else 0, + ) + attention_metadata = metadata.attention + dense_indices_cache = attention_metadata.decode_dense_compressed_indices_cache + capture_safe_keys = ( + attention_metadata.decode_dense_compressed_indices_capture_safe_keys + ) + cached = dense_indices_cache.get(cache_key) + capture_cached = cache_key in capture_safe_keys + if cached is not None and (not capturing or capture_cached): + return cached + + width = self._dense_compressed_indices_width(compress_ratio) + compressed_lens = torch.div( + positions.to(torch.int64) + 1, + compress_ratio, + rounding_mode="floor", + ).clamp(0, width) + offsets = torch.arange(width, dtype=torch.int64, device=positions.device) + local = offsets[None, :].expand(num_tokens, -1) + valid = offsets[None, :] < compressed_lens[:, None] + if is_valid_token is not None: + valid = valid & is_valid_token.to(torch.bool)[:, None] + lens = compressed_lens.to(torch.int32) + if is_valid_token is not None: + lens = torch.where( + is_valid_token.to(torch.bool), + lens, + torch.zeros_like(lens), + ) + + safe_local = torch.where(valid, local, torch.zeros_like(local)) + pages = torch.div(safe_local, block_size, rounding_mode="floor") + if block_table_base_offsets is not None: + pages = ( + pages + - block_table_base_offsets.to( + device=pages.device, + dtype=torch.int64, + )[ + req_idx + ][:, None] + ) + page_offsets = safe_local % block_size + page_ids = metadata.cache.safe_page_ids( + block_table, + req_idx[:, None], + pages.long(), + ) + slots = page_ids * block_size + page_offsets + indices_2d = torch.where( + valid & (page_ids >= 0), + slots, + torch.full_like(slots, -1), + ) + indices = indices_2d.to(torch.int32).unsqueeze(1) + dense_indices_cache[cache_key] = (indices, lens) + if capturing: + capture_safe_keys.add(cache_key) + return indices, lens + + def _dense_compressed_indices_width(self, compress_ratio: int) -> int: + if compress_ratio <= 1: + return 0 + width = max(1, (self.context_len + compress_ratio - 1) // compress_ratio) + alignment = DEEPSEEK_V4_SPARSE_PREFILL_TOPK_ALIGNMENT + return ((width + alignment - 1) // alignment) * alignment + + def _dense_prefill_local_compressed_indices( + self, + positions: torch.Tensor, + *, + compress_ratio: int, + width: int, + ) -> torch.Tensor: + shape = (positions.numel(), width) + if ( + self._prefill_dense_compressed_indices_buffer is None + or self._prefill_dense_compressed_indices_buffer.device != positions.device + or self._prefill_dense_compressed_indices_buffer.shape[0] < shape[0] + or self._prefill_dense_compressed_indices_buffer.shape[1] < shape[1] + ): + self._prefill_dense_compressed_indices_buffer = torch.empty( + shape, + dtype=torch.int32, + device=positions.device, + ) + out = self._prefill_dense_compressed_indices_buffer[: shape[0], : shape[1]] + return deepseek_v4_build_dense_prefill_local_compressed_indices( + positions=positions, + compress_ratio=compress_ratio, + width=width, + out=out, + ) + + def _get_decode_tile_metadata(self, kind: str, bs: int): + phase = ( + "graph" + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing() + else "eager" + ) + tile_metadata = self._decode_tile_metadata.get((phase, kind, bs)) + if tile_metadata is not None: + return tile_metadata + if get_mla_metadata is error_fn: + raise RuntimeError( + "DeepSeek V4 decode requires FlashMLA latent attention. " + "Build/install `tokenspeed-kernel/python` with FlashMLA." + ) + tile_metadata = get_mla_metadata()[0] + self._decode_tile_metadata[(phase, kind, bs)] = tile_metadata + return tile_metadata + + def _fp8_ds_mla_cache_view( + self, + cache_2d: torch.Tensor, + block_size: int, + ) -> torch.Tensor: + row_bytes = self._fp8_ds_mla_row_bytes + if row_bytes is None: + if cache_2d.shape[1] % block_size != 0: + raise ValueError( + "DeepSeek V4 fp8_ds_mla cache width must be divisible by " + f"block_size={block_size}, got {cache_2d.shape[1]}" + ) + row_bytes = cache_2d.shape[1] // block_size + return torch.as_strided( + cache_2d, + (cache_2d.shape[0], block_size, 1, row_bytes), + ( + cache_2d.stride(0), + row_bytes, + row_bytes, + 1, + ), + ) + + def forward_deepseek_v4_decode( + self, + *, + q: torch.Tensor, + positions: torch.Tensor, + token_to_kv_pool, + layer_id: int, + kind: str, + compress_ratio: int, + num_local_heads: int, + padded_heads: int, + head_dim: int, + window_size: int, + softmax_scale: float, + attn_sink: torch.Tensor, + topk_indices: torch.Tensor | None, + ) -> torch.Tensor: + metadata = self._select_decode_metadata(q.shape[0]) + if metadata is None: + raise RuntimeError("DeepSeek V4 decode requires forward metadata") + self.forward_metadata = metadata + if metadata.forward_mode is None or not metadata.forward_mode.is_decode(): + raise RuntimeError( + "forward_deepseek_v4_decode only supports ForwardMode.DECODE" + ) + if metadata.token_to_req_indices.numel() != q.shape[0]: + raise RuntimeError( + "DeepSeek V4 decode metadata token count mismatch: " + f"metadata_tokens={metadata.token_to_req_indices.numel()}, " + f"q_tokens={q.shape[0]}" + ) + if flash_mla_with_kvcache is error_fn: + raise RuntimeError( + "DeepSeek V4 decode requires FlashMLA latent attention. " + "Build/install `tokenspeed-kernel/python` with FlashMLA." + ) + + if q.shape[1] == padded_heads: + q_padded = q.contiguous() + else: + q_padded = torch.zeros( + (q.shape[0], padded_heads, q.shape[2]), + dtype=q.dtype, + device=q.device, + ) + q_padded[:, : q.shape[1]].copy_(q) + swa_block_size = token_to_kv_pool.swa_block_size + attention_metadata = metadata.attention + if ( + attention_metadata.decode_swa_indices is not None + and attention_metadata.decode_swa_lens is not None + and attention_metadata.decode_swa_window_size == window_size + and attention_metadata.decode_swa_block_size == swa_block_size + and attention_metadata.decode_swa_indices.shape[0] == positions.numel() + ): + swa_indices = attention_metadata.decode_swa_indices + swa_lens = attention_metadata.decode_swa_lens + else: + swa_indices, swa_lens = self._update_decode_swa_metadata( + metadata, + window_size=window_size, + block_size=swa_block_size, + ) + compressed_block_size = token_to_kv_pool.get_compressed_block_size(layer_id) + extra_indices, extra_lens = self._decode_compressed_attention_indices_and_lens( + positions, + compress_ratio=compress_ratio, + block_size=compressed_block_size, + topk_indices=topk_indices, + ) + + swa_cache = self._fp8_ds_mla_cache_view( + token_to_kv_pool.get_swa_kv_buffer(layer_id), + swa_block_size, + ) + compressed_cache = None + if compress_ratio > 1: + compressed_cache = self._fp8_ds_mla_cache_view( + token_to_kv_pool.get_compressed_kv_buffer_2d(layer_id), + compressed_block_size, + ) + + out, _ = flash_mla_with_kvcache( + q=q_padded.unsqueeze(1), + k_cache=swa_cache, + block_table=None, + cache_seqlens=None, + head_dim_v=head_dim, + tile_scheduler_metadata=self._get_decode_tile_metadata( + kind, + q_padded.shape[0], + ), + softmax_scale=softmax_scale, + is_fp8_kvcache=True, + indices=swa_indices.unsqueeze(1), + attn_sink=attn_sink, + extra_k_cache=compressed_cache, + extra_indices_in_kvcache=extra_indices, + topk_length=swa_lens, + extra_topk_length=extra_lens, + ) + if out.dim() == 4: + out = out.squeeze(1) + return out[:, :num_local_heads] + + def forward_deepseek_v4_mixed( + self, + *, + q: torch.Tensor, + positions: torch.Tensor, + token_to_kv_pool, + layer_id: int, + kind: str, + compress_ratio: int, + num_local_heads: int, + padded_heads: int, + head_dim: int, + window_size: int, + softmax_scale: float, + attn_sink: torch.Tensor, + topk_indices: torch.Tensor | None, + ) -> torch.Tensor: + metadata = self.forward_metadata + if ( + metadata is None + or metadata.forward_mode is None + or not metadata.forward_mode.is_mixed() + ): + metadata = self.forward_prefill_metadata or metadata + if ( + metadata is None + or metadata.forward_mode is None + or not metadata.forward_mode.is_mixed() + ): + raise RuntimeError("DeepSeek V4 mixed attention requires forward metadata") + + num_prefill_reqs = metadata.num_prefill_reqs + num_prefill_tokens = metadata.num_prefill_tokens + num_decode_reqs = metadata.decode_req_count() + num_decode_tokens = metadata.decode_token_count() + out = q.new_empty((q.shape[0], num_local_heads, head_dim)) + saved_metadata = self.forward_metadata + try: + if num_prefill_tokens > 0: + self.forward_metadata = self._metadata_slice( + metadata, + req_start=0, + req_end=num_prefill_reqs, + token_start=0, + token_end=num_prefill_tokens, + forward_mode=ForwardMode.EXTEND, + ) + prefill_out = self.forward_deepseek_v4_prefill( + q=q[:num_prefill_tokens], + positions=positions[:num_prefill_tokens], + token_to_kv_pool=token_to_kv_pool, + layer_id=layer_id, + kind=kind, + compress_ratio=compress_ratio, + num_local_heads=num_local_heads, + padded_heads=padded_heads, + head_dim=head_dim, + window_size=window_size, + softmax_scale=softmax_scale, + attn_sink=attn_sink, + topk_indices=( + topk_indices[:num_prefill_tokens] + if topk_indices is not None + else None + ), + ) + with nvtx_range(f"attn_{kind}_mixed_prefill_copy"): + out[:num_prefill_tokens].copy_(prefill_out) + if num_decode_tokens > 0: + decode_end = num_prefill_tokens + num_decode_tokens + self.forward_metadata = self._metadata_slice( + metadata, + req_start=num_prefill_reqs, + req_end=num_prefill_reqs + num_decode_reqs, + token_start=num_prefill_tokens, + token_end=decode_end, + forward_mode=ForwardMode.DECODE, + ) + decode_out = self.forward_deepseek_v4_decode( + q=q[num_prefill_tokens:decode_end], + positions=positions[num_prefill_tokens:decode_end], + token_to_kv_pool=token_to_kv_pool, + layer_id=layer_id, + kind=kind, + compress_ratio=compress_ratio, + num_local_heads=num_local_heads, + padded_heads=padded_heads, + head_dim=head_dim, + window_size=window_size, + softmax_scale=softmax_scale, + attn_sink=attn_sink, + topk_indices=( + topk_indices[num_prefill_tokens:decode_end] + if topk_indices is not None + else None + ), + ) + with nvtx_range(f"attn_{kind}_mixed_decode_copy"): + out[num_prefill_tokens:decode_end].copy_(decode_out) + finally: + self.forward_metadata = saved_metadata + return out + + def _prefill_workspace( + self, + *, + positions: torch.Tensor, + token_to_kv_pool, + layer_id: int, + compress_ratio: int, + window_size: int, + head_dim: int, + topk_indices: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + metadata = self.forward_metadata + if metadata is None: + raise RuntimeError("DeepSeek V4 prefill requires forward metadata") + cache_metadata = metadata.cache + num_reqs = metadata.seq_lens.numel() + prefix_lens = metadata.seq_lens - metadata.query_lens + gather_lens = metadata.query_lens + torch.minimum( + prefix_lens, + torch.full_like(prefix_lens, max(window_size - 1, 0)), + ) + if cache_metadata.swa_block_table is None: + raise RuntimeError("DeepSeek V4 missing paged-cache block table for SWA KV") + swa_block_table = cache_metadata.swa_block_table + max_gather_len = int(gather_lens.max().item()) if num_reqs else 1 + compressed_lens = ( + torch.div(metadata.seq_lens, compress_ratio, rounding_mode="floor") + if compress_ratio > 1 + else torch.zeros_like(metadata.seq_lens) + ) + compressed_base = ( + int(compressed_lens.max().item()) if compress_ratio > 1 and num_reqs else 0 + ) + workspace_width = max(1, compressed_base + max_gather_len) + kv_workspace = self._get_prefill_workspace( + num_reqs=num_reqs, + workspace_width=workspace_width, + head_dim=head_dim, + device=positions.device, + ) + + if compress_ratio == 4 and topk_indices is not None: + compressed_block_size = token_to_kv_pool.get_compressed_block_size(layer_id) + compressed_cache = token_to_kv_pool.get_compressed_kv_buffer_2d(layer_id) + compressed_block_table = cache_metadata.compressed_block_table( + compress_ratio, + compressed_block_size, + ) + deepseek_v4_dequantize_and_gather_k_cache( + out=kv_workspace, + cache_2d=compressed_cache, + seq_lens=compressed_lens, + gather_lens=None, + block_table=compressed_block_table, + block_size=compressed_block_size, + offset=0, + ) + deepseek_v4_dequantize_and_gather_k_cache( + out=kv_workspace, + cache_2d=token_to_kv_pool.get_swa_kv_buffer(layer_id), + seq_lens=metadata.seq_lens, + gather_lens=gather_lens, + block_table=swa_block_table, + block_table_base_offsets=cache_metadata.swa_base_logical_page, + block_size=token_to_kv_pool.swa_block_size, + offset=compressed_base, + ) + indices, lens = deepseek_v4_combine_topk_swa_indices( + topk_indices=topk_indices, + query_start_loc=metadata.query_start_loc, + seq_lens=metadata.seq_lens, + gather_lens=gather_lens, + window_size=window_size, + compress_ratio=compress_ratio, + topk=topk_indices.shape[-1], + workspace_width=workspace_width, + compressed_base=compressed_base, + ) + return kv_workspace, indices, lens + + if compress_ratio == 4: + raise RuntimeError("DeepSeek V4 CSA prefill requires top-k indices") + + swa_cache = token_to_kv_pool.get_swa_kv_buffer(layer_id) + compressed_cache = ( + token_to_kv_pool.get_compressed_kv_buffer_2d(layer_id) + if compress_ratio > 1 + else None + ) + if compress_ratio > 1: + assert compressed_cache is not None + compressed_block_size = token_to_kv_pool.get_compressed_block_size(layer_id) + compressed_block_table = cache_metadata.compressed_block_table( + compress_ratio, + compressed_block_size, + ) + deepseek_v4_dequantize_and_gather_k_cache( + out=kv_workspace, + cache_2d=compressed_cache, + seq_lens=compressed_lens, + gather_lens=None, + block_table=compressed_block_table, + block_size=compressed_block_size, + offset=0, + ) + deepseek_v4_dequantize_and_gather_k_cache( + out=kv_workspace, + cache_2d=swa_cache, + seq_lens=metadata.seq_lens, + gather_lens=gather_lens, + block_table=swa_block_table, + block_table_base_offsets=cache_metadata.swa_base_logical_page, + block_size=token_to_kv_pool.swa_block_size, + offset=compressed_base, + ) + if compress_ratio > 1: + dense_compressed_indices = self._dense_prefill_local_compressed_indices( + positions, + compress_ratio=compress_ratio, + width=self._dense_compressed_indices_width(compress_ratio), + ) + indices, lens = deepseek_v4_combine_topk_swa_indices( + topk_indices=dense_compressed_indices, + query_start_loc=metadata.query_start_loc, + seq_lens=metadata.seq_lens, + gather_lens=gather_lens, + window_size=window_size, + compress_ratio=compress_ratio, + topk=dense_compressed_indices.shape[-1], + workspace_width=workspace_width, + compressed_base=compressed_base, + ) + return kv_workspace, indices, lens + + indices, lens = deepseek_v4_combine_dense_swa_indices( + positions=positions, + token_to_req_indices=metadata.token_to_req_indices[: positions.numel()], + seq_lens=metadata.seq_lens, + compressed_lens=compressed_lens, + gather_lens=gather_lens, + window_size=window_size, + compress_ratio=compress_ratio, + workspace_width=workspace_width, + compressed_base=compressed_base, + ) + return kv_workspace, indices, lens + + def _metadata_slice( + self, + metadata: DeepseekV4ForwardMetadata, + *, + req_start: int, + req_end: int, + token_start: int, + token_end: int, + forward_mode: ForwardMode, + ) -> DeepseekV4ForwardMetadata: + token_to_req = metadata.token_to_req_indices[token_start:token_end].to( + torch.int32 + ) - int(req_start) + cache_metadata = metadata.cache + paged_cache_block_tables = { + key: table[req_start:req_end] + for key, table in cache_metadata.paged_cache_block_tables.items() + } + paged_cache_block_table_base_offsets = { + key: offsets[req_start:req_end] + for key, offsets in ( + cache_metadata.paged_cache_block_table_base_offsets.items() + ) + } + compressor_state_block_tables = { + key: table[req_start:req_end] + for key, table in cache_metadata.compressor_state_block_tables.items() + } + compressor_state_base_logical_pages = { + key: offsets[req_start:req_end] + for key, offsets in ( + cache_metadata.compressor_state_base_logical_pages.items() + ) + } + query_lens = metadata.query_lens[req_start:req_end] + req_count = max(0, req_end - req_start) + token_count = max(0, token_end - token_start) + num_prefill_reqs = req_count if forward_mode.is_extend_or_mixed() else 0 + num_prefill_tokens = token_count if forward_mode.is_extend_or_mixed() else 0 + query_start_loc = torch.nn.functional.pad( + torch.cumsum(query_lens.to(torch.int32), dim=0, dtype=torch.int32), + (1, 0), + ) + sliced_cache = DeepseekV4CacheMetadata( + page_size=cache_metadata.page_size, + block_table=cache_metadata.block_table[req_start:req_end], + paged_cache_block_tables=paged_cache_block_tables, + paged_cache_block_table_base_offsets=paged_cache_block_table_base_offsets, + swa_block_table=( + cache_metadata.swa_block_table[req_start:req_end] + if cache_metadata.swa_block_table is not None + else None + ), + swa_base_logical_page=( + cache_metadata.swa_base_logical_page[req_start:req_end] + if cache_metadata.swa_base_logical_page is not None + else None + ), + compressor_state_block_tables=compressor_state_block_tables, + compressor_state_base_logical_pages=compressor_state_base_logical_pages, + indexer_state_block_table=( + cache_metadata.indexer_state_block_table[req_start:req_end] + if cache_metadata.indexer_state_block_table is not None + else None + ), + indexer_state_base_logical_page=( + cache_metadata.indexer_state_base_logical_page[req_start:req_end] + if cache_metadata.indexer_state_base_logical_page is not None + else None + ), + ) + return DeepseekV4ForwardMetadata( + req_pool_indices=metadata.req_pool_indices[req_start:req_end], + seq_lens=metadata.seq_lens[req_start:req_end], + query_lens=query_lens, + query_start_loc=query_start_loc, + token_to_req_indices=token_to_req, + cache=sliced_cache, + is_valid_token=( + metadata.is_valid_token[token_start:token_end] + if metadata.is_valid_token is not None + else None + ), + seq_lens_cpu=( + metadata.seq_lens_cpu[req_start:req_end] + if metadata.seq_lens_cpu is not None + else None + ), + query_lens_cpu=( + metadata.query_lens_cpu[req_start:req_end] + if metadata.query_lens_cpu is not None + else None + ), + num_prefill_reqs=num_prefill_reqs, + num_prefill_tokens=num_prefill_tokens, + forward_mode=forward_mode, + ) + + def _forward_deepseek_v4_prefill_chunk( + self, + *, + q: torch.Tensor, + positions: torch.Tensor, + token_to_kv_pool, + layer_id: int, + kind: str, + compress_ratio: int, + num_local_heads: int, + padded_heads: int, + head_dim: int, + window_size: int, + softmax_scale: float, + attn_sink: torch.Tensor, + topk_indices: torch.Tensor | None, + ) -> torch.Tensor: + metadata = self.forward_metadata + if metadata is None: + raise RuntimeError("DeepSeek V4 prefill requires forward metadata") + if flash_mla_sparse_fwd is error_fn: + raise RuntimeError( + "DeepSeek V4 prefill requires FlashMLA sparse attention. " + "Build/install `tokenspeed-kernel/python` with FlashMLA." + ) + + with nvtx_range(f"attn_{kind}_prefill_pad_q"): + if q.shape[1] == padded_heads: + q_padded = q.contiguous() + else: + q_padded = torch.zeros( + (q.shape[0], padded_heads, q.shape[2]), + dtype=q.dtype, + device=q.device, + ) + q_padded[:, : q.shape[1]].copy_(q) + with nvtx_range(f"attn_{kind}_prefill_workspace"): + kv_workspace, indices, lens = self._prefill_workspace( + positions=positions, + token_to_kv_pool=token_to_kv_pool, + layer_id=layer_id, + compress_ratio=compress_ratio, + window_size=window_size, + head_dim=head_dim, + topk_indices=topk_indices, + ) + with nvtx_range(f"attn_{kind}_prefill_flashmla"): + out, _, _ = flash_mla_sparse_fwd( + q=q_padded, + kv=kv_workspace.view(-1, 1, head_dim), + indices=indices.unsqueeze(1), + sm_scale=softmax_scale, + attn_sink=attn_sink, + topk_length=lens, + ) + return out[:, :num_local_heads] + + def forward_deepseek_v4_prefill( + self, + *, + q: torch.Tensor, + positions: torch.Tensor, + token_to_kv_pool, + layer_id: int, + kind: str, + compress_ratio: int, + num_local_heads: int, + padded_heads: int, + head_dim: int, + window_size: int, + softmax_scale: float, + attn_sink: torch.Tensor, + topk_indices: torch.Tensor | None, + ) -> torch.Tensor: + metadata = self.forward_metadata + if ( + metadata is None + or metadata.forward_mode is None + or not metadata.forward_mode.is_extend_or_mixed() + ): + metadata = self.forward_prefill_metadata or metadata + if metadata is None: + raise RuntimeError("DeepSeek V4 prefill requires forward metadata") + self.forward_metadata = metadata + if ( + metadata.forward_mode is None + or not metadata.forward_mode.is_extend_or_mixed() + ): + raise RuntimeError( + "forward_deepseek_v4_prefill only supports extend/prefill modes" + ) + if metadata.token_to_req_indices.numel() != q.shape[0]: + raise RuntimeError( + "DeepSeek V4 prefill metadata token count mismatch: " + f"metadata_tokens={metadata.token_to_req_indices.numel()}, " + f"q_tokens={q.shape[0]}" + ) + + num_reqs = int(metadata.num_prefill_reqs or metadata.seq_lens.numel()) + if num_reqs <= self.prefill_chunk_size: + return self._forward_deepseek_v4_prefill_chunk( + q=q, + positions=positions, + token_to_kv_pool=token_to_kv_pool, + layer_id=layer_id, + kind=kind, + compress_ratio=compress_ratio, + num_local_heads=num_local_heads, + padded_heads=padded_heads, + head_dim=head_dim, + window_size=window_size, + softmax_scale=softmax_scale, + attn_sink=attn_sink, + topk_indices=topk_indices, + ) + + token_offsets = [ + int(x) + for x in metadata.query_start_loc[: num_reqs + 1].detach().cpu().tolist() + ] + out = q.new_empty((q.shape[0], num_local_heads, head_dim)) + saved_metadata = self.forward_metadata + try: + for req_start in range(0, num_reqs, self.prefill_chunk_size): + req_end = min(req_start + self.prefill_chunk_size, num_reqs) + token_start = token_offsets[req_start] + token_end = token_offsets[req_end] + if token_end <= token_start: + continue + self.forward_metadata = self._metadata_slice( + saved_metadata, + req_start=req_start, + req_end=req_end, + token_start=token_start, + token_end=token_end, + forward_mode=ForwardMode.EXTEND, + ) + chunk_out = self._forward_deepseek_v4_prefill_chunk( + q=q[token_start:token_end], + positions=positions[token_start:token_end], + token_to_kv_pool=token_to_kv_pool, + layer_id=layer_id, + kind=kind, + compress_ratio=compress_ratio, + num_local_heads=num_local_heads, + padded_heads=padded_heads, + head_dim=head_dim, + window_size=window_size, + softmax_scale=softmax_scale, + attn_sink=attn_sink, + topk_indices=( + topk_indices[token_start:token_end] + if topk_indices is not None + else None + ), + ) + out[token_start:token_end].copy_(chunk_out) + finally: + self.forward_metadata = saved_metadata + return out + + def init_cuda_graph_state( + self, + max_bs: int, + seq_lens_buf: torch.Tensor | None = None, + paged_cache_group_specs=(), + max_tokens_per_req: int = 1, + overlap_schedule_depth: int = 0, + ): + del seq_lens_buf + self._decode_tile_metadata = {} + self._cuda_graph_max_tokens_per_req = max( + 1, + int(max_tokens_per_req), + int(self.speculative_num_draft_tokens or 0), + ) + max_tokens = max_bs * self._cuda_graph_max_tokens_per_req + self._cuda_graph_block_table = torch.zeros( + (max_bs, self.max_num_pages), + dtype=torch.int32, + device=self.device, + ) + self._cuda_graph_req_pool_indices = torch.zeros( + (max_bs,), + dtype=torch.int32, + device=self.device, + ) + self._cuda_graph_seq_lens = torch.ones( + (max_bs,), + dtype=torch.int32, + device=self.device, + ) + self._cuda_graph_query_lens = torch.ones( + (max_bs,), + dtype=torch.int32, + device=self.device, + ) + self._cuda_graph_query_start_loc = torch.arange( + max_bs + 1, + dtype=torch.int32, + device=self.device, + ) + self._cuda_graph_token_to_req = torch.arange( + max_tokens, + dtype=torch.int32, + device=self.device, + ) + query_start_base = torch.arange( + max_bs + 1, + dtype=torch.int32, + device=self.device, + ) + token_to_req_base = torch.arange( + max_bs, + dtype=torch.int32, + device=self.device, + ) + self._cuda_graph_query_start_by_tokens_per_req = {} + self._cuda_graph_token_to_req_by_tokens_per_req = {} + for tokens_per_req in range(1, self._cuda_graph_max_tokens_per_req + 1): + self._cuda_graph_query_start_by_tokens_per_req[tokens_per_req] = ( + query_start_base * tokens_per_req + ) + self._cuda_graph_token_to_req_by_tokens_per_req[tokens_per_req] = ( + token_to_req_base.repeat_interleave(tokens_per_req) + ) + self._cuda_graph_max_bs = max_bs + self._cuda_graph_paged_cache_block_tables = {} + self._cuda_graph_paged_cache_base_offsets = {} + for spec in tuple(paged_cache_group_specs or ()): + gid = str(spec.group_id) + sliding = str(getattr(spec, "retention", "")) == "sliding_window" + max_pages = compute_max_logical_pages_for_capture( + spec, + max_context_len=self.context_len, + max_tokens_per_req=max_tokens_per_req, + overlap_schedule_depth=overlap_schedule_depth, + ) + self._cuda_graph_paged_cache_block_tables[gid] = torch.zeros( + (max_bs, max_pages), + dtype=torch.int32, + device=self.device, + ) + if sliding: + self._cuda_graph_paged_cache_base_offsets[gid] = torch.zeros( + (max_bs,), + dtype=torch.int32, + device=self.device, + ) + self._cuda_graph_is_valid_token = torch.ones( + max_tokens, + dtype=torch.bool, + device=self.device, + ) + + def _refresh_cuda_graph_paged_cache_block_tables( + self, + bs: int, + paged_cache_block_tables: dict[str, torch.Tensor], + *, + pad_value: int, + ) -> dict[str, torch.Tensor]: + out: dict[str, torch.Tensor] = {} + if not self._cuda_graph_paged_cache_block_tables: + return out + for group_id, buf in self._cuda_graph_paged_cache_block_tables.items(): + table = paged_cache_block_tables.get(group_id) + buf[:bs].fill_(pad_value) + if table is not None: + if int(table.shape[0]) != bs: + raise RuntimeError( + "DeepSeek V4 CUDA graph paged cache table row count " + f"mismatch for {group_id!r}: got {int(table.shape[0])}, " + f"expected padded bs {bs}" + ) + cols = int(table.shape[1]) + if cols > int(buf.shape[1]): + raise RuntimeError( + "DeepSeek V4 CUDA graph paged cache table width " + f"mismatch for {group_id!r}: got {cols}, capture " + f"buffer has {int(buf.shape[1])}" + ) + if cols > 0: + buf[:bs, :cols].copy_(table[:bs, :cols].to(torch.int32)) + out[group_id] = buf[:bs] + return out + + def _refresh_cuda_graph_base_offsets( + self, + bs: int, + base_offsets: dict[str, torch.Tensor], + ) -> dict[str, torch.Tensor]: + """Refresh persistent base-offset buffers from per-step input. + + Sliding groups whose key is missing fall back to 0. Returns the [:bs] + views keyed by gid. + """ + out: dict[str, torch.Tensor] = {} + for gid, buf in self._cuda_graph_paged_cache_base_offsets.items(): + buf[:bs].fill_(0) + src = base_offsets.get(gid) + if src is not None and bs > 0: + rows = int(src.shape[0]) + if rows < bs: + raise RuntimeError( + "DeepSeek V4 CUDA-graph replay base-offsets row count " + f"{rows} < bs={bs} for group {gid!r}" + ) + buf[:bs].copy_(src[:bs].to(torch.int32)) + out[gid] = buf[:bs] + return out + + def _cuda_graph_tokens_per_req( + self, + bs: int, + num_tokens: int, + forward_mode: ForwardMode | None, + ) -> int: + if num_tokens != bs: + if bs == 0: + return self._cuda_graph_max_tokens_per_req + if num_tokens % bs != 0: + raise RuntimeError( + "DeepSeek V4 packed CUDA graph metadata expects uniformly " + f"packed tokens per request, got num_tokens={num_tokens}, " + f"bs={bs}" + ) + tokens_per_req = num_tokens // bs + if tokens_per_req > self._cuda_graph_max_tokens_per_req: + raise RuntimeError( + "DeepSeek V4 packed CUDA graph metadata was initialized " + f"for at most {self._cuda_graph_max_tokens_per_req} tokens " + f"per request, got {tokens_per_req}" + ) + return max(1, tokens_per_req) + return 1 + + def _refresh_cuda_graph_packed_metadata( + self, + *, + bs: int, + actual_bs: int, + tokens_per_req: int, + ) -> int: + total_tokens = bs * tokens_per_req + actual_tokens = actual_bs * tokens_per_req + query_start = self._cuda_graph_query_start_by_tokens_per_req.get(tokens_per_req) + token_to_req = self._cuda_graph_token_to_req_by_tokens_per_req.get( + tokens_per_req + ) + if query_start is None or token_to_req is None: + raise RuntimeError( + "DeepSeek V4 CUDA graph packed metadata was not precomputed " + f"for tokens_per_req={tokens_per_req}" + ) + self._cuda_graph_query_lens[:bs].fill_(tokens_per_req) + self._cuda_graph_query_start_loc[: bs + 1].copy_(query_start[: bs + 1]) + self._cuda_graph_token_to_req[:total_tokens].copy_(token_to_req[:total_tokens]) + self._cuda_graph_is_valid_token[:actual_tokens].fill_(True) + if actual_tokens < total_tokens: + self._cuda_graph_is_valid_token[actual_tokens:total_tokens].fill_(False) + return total_tokens + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + forward_mode: ForwardMode, + **kwargs, + ): + paged_cache_block_tables = kwargs.pop("paged_cache_block_tables", None) or {} + paged_cache_block_table_base_offsets = ( + kwargs.pop("paged_cache_block_table_base_offsets", None) or {} + ) + num_tokens_arg = kwargs.pop("num_tokens", None) + del kwargs + if forward_mode is not None and not forward_mode.is_decode_or_idle(): + raise NotImplementedError( + f"DeepSeek V4 CUDA graph capture not supported for {forward_mode}" + ) + if num_tokens_arg is None: + num_tokens = bs + else: + num_tokens = int(num_tokens_arg) + tokens_per_req = self._cuda_graph_tokens_per_req(bs, num_tokens, forward_mode) + is_packed_decode = ( + forward_mode is not None and forward_mode.is_decode() and num_tokens != bs + ) + total_tokens = self._refresh_cuda_graph_packed_metadata( + bs=bs, + actual_bs=bs, + tokens_per_req=tokens_per_req, + ) + capture_seq_lens = seq_lens[:bs].to(torch.int32) + if is_packed_decode: + capture_seq_lens = torch.maximum( + capture_seq_lens, + torch.full_like(capture_seq_lens, tokens_per_req), + ) + self._cuda_graph_req_pool_indices[:bs].copy_(req_pool_indices[:bs]) + self._cuda_graph_seq_lens[:bs].copy_(capture_seq_lens) + metadata_forward_mode = forward_mode + is_decode = ( + metadata_forward_mode is not None and metadata_forward_mode.is_decode() + ) + offsets_on_device = { + str(gid): off.to(device=self.device, dtype=torch.int32) + for gid, off in paged_cache_block_table_base_offsets.items() + } + metadata_paged = self._refresh_cuda_graph_paged_cache_block_tables( + bs, + { + str(group_id): table.to(device=self.device, dtype=torch.int32) + for group_id, table in paged_cache_block_tables.items() + }, + pad_value=0, + ) + metadata_base_offsets = self._refresh_cuda_graph_base_offsets( + bs, + offsets_on_device, + ) + ( + swa_block_table, + compressor_state_block_tables, + indexer_state_block_table, + swa_base, + compressor_state_base, + indexer_state_base, + ) = _split_paged_cache_block_tables_into_v4_metadata( + metadata_paged, + metadata_base_offsets, + ) + prior_metadata = self._cuda_graph_metadata.get(bs) + prior_slot_mappings = ( + prior_metadata.cache.decode_compressed_slot_mappings + if prior_metadata is not None + else {} + ) + cache_metadata = DeepseekV4CacheMetadata( + page_size=self.page_size, + block_table=self._cuda_graph_block_table[:bs, : self.max_num_pages], + paged_cache_block_tables=metadata_paged, + paged_cache_block_table_base_offsets=metadata_base_offsets, + swa_block_table=swa_block_table, + swa_base_logical_page=swa_base, + compressor_state_block_tables=compressor_state_block_tables, + compressor_state_base_logical_pages=compressor_state_base, + indexer_state_block_table=indexer_state_block_table, + indexer_state_base_logical_page=indexer_state_base, + decode_compressed_slot_mappings=prior_slot_mappings, + ) + metadata = prior_metadata + if metadata is None: + metadata = DeepseekV4ForwardMetadata( + req_pool_indices=self._cuda_graph_req_pool_indices[:bs], + seq_lens=self._cuda_graph_seq_lens[:bs], + query_lens=self._cuda_graph_query_lens[:bs], + query_start_loc=self._cuda_graph_query_start_loc[: bs + 1], + token_to_req_indices=self._cuda_graph_token_to_req[:total_tokens], + cache=cache_metadata, + is_valid_token=self._cuda_graph_is_valid_token[:total_tokens], + seq_lens_cpu=None, + query_lens_cpu=None, + forward_mode=metadata_forward_mode, + ) + else: + metadata.req_pool_indices = self._cuda_graph_req_pool_indices[:bs] + metadata.seq_lens = self._cuda_graph_seq_lens[:bs] + metadata.query_lens = self._cuda_graph_query_lens[:bs] + metadata.query_start_loc = self._cuda_graph_query_start_loc[: bs + 1] + metadata.token_to_req_indices = self._cuda_graph_token_to_req[:total_tokens] + metadata.cache = cache_metadata + metadata.is_valid_token = self._cuda_graph_is_valid_token[:total_tokens] + metadata.seq_lens_cpu = None + metadata.query_lens_cpu = None + metadata.forward_mode = metadata_forward_mode + self._cuda_graph_metadata[bs] = metadata + if is_packed_decode and getattr(self, "is_draft", False): + self._prepare_draft_decode_metadata( + metadata, + self._cuda_graph_seq_lens[:bs], + ) + if is_decode: + self.forward_decode_metadata = metadata + self.forward_metadata = metadata + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + forward_mode: ForwardMode = None, + req_to_page: torch.Tensor = None, + **kwargs, + ): + paged_cache_block_tables = kwargs.pop("paged_cache_block_tables", None) or {} + paged_cache_block_table_base_offsets = ( + kwargs.pop("paged_cache_block_table_base_offsets", None) or {} + ) + actual_bs = max(0, min(int(kwargs.pop("actual_bs", bs)), bs)) + num_tokens_arg = kwargs.pop("num_tokens", None) + del kwargs + if forward_mode is not None and not forward_mode.is_decode_or_idle(): + raise NotImplementedError( + f"DeepSeek V4 CUDA graph replay not supported for {forward_mode}" + ) + if num_tokens_arg is None: + num_tokens = bs + else: + num_tokens = int(num_tokens_arg) + tokens_per_req = self._cuda_graph_tokens_per_req(bs, num_tokens, forward_mode) + is_packed_decode = ( + forward_mode is not None and forward_mode.is_decode() and num_tokens != bs + ) + total_tokens = self._refresh_cuda_graph_packed_metadata( + bs=bs, + actual_bs=actual_bs, + tokens_per_req=tokens_per_req, + ) + metadata = self._cuda_graph_metadata[bs] + self._cuda_graph_req_pool_indices[:bs].copy_(req_pool_indices[:bs]) + self._cuda_graph_seq_lens[:bs].copy_(seq_lens[:bs].to(torch.int32)) + if req_to_page is not None: + self._cuda_graph_block_table[:bs, : self.max_num_pages].copy_( + req_to_page[req_pool_indices[:bs], : self.max_num_pages] + ) + offsets_on_device = { + str(gid): off.to(device=self.device, dtype=torch.int32) + for gid, off in paged_cache_block_table_base_offsets.items() + } + metadata_paged = self._refresh_cuda_graph_paged_cache_block_tables( + bs, + { + str(group_id): table.to(device=self.device, dtype=torch.int32) + for group_id, table in paged_cache_block_tables.items() + }, + pad_value=-1, + ) + metadata_base_offsets = self._refresh_cuda_graph_base_offsets( + bs, + offsets_on_device, + ) + ( + swa_block_table, + compressor_state_block_tables, + indexer_state_block_table, + swa_base, + compressor_state_base, + indexer_state_base, + ) = _split_paged_cache_block_tables_into_v4_metadata( + metadata_paged, + metadata_base_offsets, + ) + metadata_forward_mode = forward_mode + is_decode = ( + metadata_forward_mode is not None and metadata_forward_mode.is_decode() + ) + metadata.forward_mode = metadata_forward_mode + metadata.token_to_req_indices = self._cuda_graph_token_to_req[:total_tokens] + metadata.is_valid_token = self._cuda_graph_is_valid_token[:total_tokens] + metadata.cache = DeepseekV4CacheMetadata( + page_size=self.page_size, + block_table=self._cuda_graph_block_table[:bs, : self.max_num_pages], + paged_cache_block_tables=metadata_paged, + paged_cache_block_table_base_offsets=metadata_base_offsets, + swa_block_table=swa_block_table, + swa_base_logical_page=swa_base, + compressor_state_block_tables=compressor_state_block_tables, + compressor_state_base_logical_pages=compressor_state_base, + indexer_state_block_table=indexer_state_block_table, + indexer_state_base_logical_page=indexer_state_base, + decode_compressed_slot_mappings=( + metadata.cache.decode_compressed_slot_mappings + ), + ) + metadata.num_prefill_reqs = 0 + metadata.num_prefill_tokens = 0 + if is_packed_decode and getattr(self, "is_draft", False): + self._prepare_draft_decode_metadata( + metadata, + self._cuda_graph_seq_lens[:bs], + ) + if ( + metadata_forward_mode is not None + and metadata_forward_mode.is_decode() + and self._decode_swa_window_size > 0 + and self._decode_swa_block_size > 0 + ): + self._update_decode_swa_metadata( + metadata, + window_size=self._decode_swa_window_size, + block_size=self._decode_swa_block_size, + ) + metadata.cache.refresh_decode_compressed_slot_mappings( + token_to_req_indices=metadata.token_to_req_indices, + query_start_loc=metadata.query_start_loc, + seq_lens=metadata.seq_lens, + is_valid_token=metadata.is_valid_token, + ) + _refresh_decode_indexer_plan_cache( + metadata, + max_context_len=self.context_len, + ) + _refresh_decode_indexer_schedule_metadata(metadata) + if is_decode: + self.forward_decode_metadata = metadata + self.forward_metadata = metadata + + def advance_draft_forward_metadata(self, seq_lens: torch.Tensor | None = None): + if ( + self._draft_decode_base_seq_lens is None + or self.forward_prefill_metadata is None + or self._draft_decode_metadata is None + ): + raise RuntimeError("DeepSeek V4 draft metadata was not initialized") + self._draft_decode_step += 1 + metadata = self._draft_decode_metadata + if seq_lens is None: + metadata.seq_lens.add_(1) + else: + metadata.seq_lens.copy_(seq_lens[: metadata.seq_lens.numel()]) + metadata.forward_mode = ForwardMode.DECODE + if self._decode_swa_window_size > 0 and self._decode_swa_block_size > 0: + self._update_decode_swa_metadata( + metadata, + window_size=self._decode_swa_window_size, + block_size=self._decode_swa_block_size, + ) + # seq_lens just changed, so any previously-refreshed plan tensors are + # stale. Re-run the same metadata-setup hooks the main path uses. + metadata.cache.refresh_decode_compressed_slot_mappings( + token_to_req_indices=metadata.token_to_req_indices, + query_start_loc=metadata.query_start_loc, + seq_lens=metadata.seq_lens, + is_valid_token=metadata.is_valid_token, + ) + _refresh_decode_indexer_plan_cache( + metadata, + max_context_len=self.context_len, + ) + _refresh_decode_indexer_schedule_metadata(metadata) + self.forward_decode_metadata = metadata + self.forward_metadata = metadata + self._decode_tile_metadata = {} + + def forward_decode(self, *args, **kwargs): + raise NotImplementedError("DeepSeek V4 uses the model-local attention forward") + + def forward_extend(self, *args, **kwargs): + raise NotImplementedError("DeepSeek V4 uses the model-local attention forward") + + +register_backend("deepseek_v4", {AttentionArch.MLA}, DeepseekV4AttentionBackend) diff --git a/python/tokenspeed/runtime/layers/attention/backends/dsa.py b/python/tokenspeed/runtime/layers/attention/backends/dsa.py new file mode 100644 index 0000000..c9e2321 --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/backends/dsa.py @@ -0,0 +1,569 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import torch +from tokenspeed_kernel.ops.attention import ( + dsa_decode, + dsa_plan, + dsa_prefill, +) +from tokenspeed_kernel.ops.attention.triton.dsa_topk import ( + workspace_topk_to_global_slots, +) +from tokenspeed_kernel.platform import current_platform + +from tokenspeed.runtime.configs.model_config import AttentionArch +from tokenspeed.runtime.execution.forward_batch_info import ForwardMode +from tokenspeed.runtime.layers.attention.backends.base import AttentionBackend +from tokenspeed.runtime.layers.attention.backends.mla import MLAAttnBackend +from tokenspeed.runtime.layers.attention.backends.trtllm_mla import TRTLLMMLABackend +from tokenspeed.runtime.layers.attention.configs.dsa import DSAConfig +from tokenspeed.runtime.layers.attention.registry import register_backend + + +def _make_dense_backend(config: DSAConfig, platform) -> AttentionBackend: + if platform.is_nvidia: + return TRTLLMMLABackend(config) + if platform.is_amd: + return MLAAttnBackend(config) + raise RuntimeError(f"DSA backend does not support platform {platform.vendor!r}.") + + +class DSABackend(AttentionBackend): + """DSA backend for sparse MLA attention. + + Dense MLA metadata and dense attention calls are delegated to a platform backend. + """ + + def __init__(self, config: DSAConfig): + super().__init__(config) + platform = current_platform() + self._dense_backend = _make_dense_backend(config, platform) + self.index_topk = config.index_topk + self.max_context_len = config.context_len + self.page_size = config.page_size + 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 = config.kv_cache_dim + self.scaling = config.scaling + self.data_type = config.kv_cache_dtype + self.q_data_type = config.dtype + self.num_local_heads = config.num_attention_heads // config.attn_tp_size + self._prefill_block_tables: torch.Tensor | None = None + + @property + def forward_decode_metadata(self): + return self._dense_backend.forward_decode_metadata + + @property + def forward_prefill_metadata(self): + return self._dense_backend.forward_prefill_metadata + + @property + def chunked_prefill_metadata(self): + return self._dense_backend.chunked_prefill_metadata + + @property + def decode_cuda_graph_metadata(self): + return self._dense_backend.decode_cuda_graph_metadata + + @property + def decode_cuda_graph_kv_indices(self): + return getattr(self._dense_backend, "decode_cuda_graph_kv_indices", None) + + @decode_cuda_graph_kv_indices.setter + def decode_cuda_graph_kv_indices(self, value): + if not hasattr(self._dense_backend, "decode_cuda_graph_kv_indices"): + raise RuntimeError( + "DSA dense backend does not expose decode CUDA graph KV indices." + ) + self._dense_backend.decode_cuda_graph_kv_indices = value + + @property + def trtllm_workspace(self): + return self._dense_backend.trtllm_workspace + + @property + def _block_table_aliased(self): + return getattr(self._dense_backend, "_block_table_aliased", False) + + @_block_table_aliased.setter + def _block_table_aliased(self, value): + if hasattr(self, "_dense_backend"): + self._dense_backend._block_table_aliased = value + + def register_step_counter(self, step_counter): + super().register_step_counter(step_counter) + self._dense_backend.register_step_counter(step_counter) + + def override_num_extends(self, num_extends: int): + return self._dense_backend.override_num_extends(num_extends) + + def init_cuda_graph_state(self, max_bs: int, seq_lens_buf: torch.Tensor): + self._dense_backend.init_cuda_graph_state(max_bs, seq_lens_buf) + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + forward_mode: ForwardMode, + ): + self._dense_backend.init_forward_metadata_capture_cuda_graph( + bs=bs, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + forward_mode=forward_mode, + ) + metadata = self.forward_decode_metadata + # Full-length broadcast: the plan and paged-MQA-logits kernels read only + # the last column, and the per-token causal bound is applied downstream. + metadata._dsa_seq_lens_2d = ( + seq_lens.unsqueeze(1).expand(-1, self.spec_num_tokens).contiguous() + ) + metadata._dsa_plan = dsa_plan( + seq_lens_2d=metadata._dsa_seq_lens_2d, page_size=self.page_size + ) + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + forward_mode: ForwardMode = None, + req_to_page: torch.Tensor = None, + **kwargs, + ): + self._dense_backend.init_forward_metadata_replay_cuda_graph( + bs=bs, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + forward_mode=forward_mode, + req_to_page=req_to_page, + **kwargs, + ) + metadata = self.forward_decode_metadata + metadata._dsa_seq_lens_2d.copy_( + seq_lens.unsqueeze(1).expand(-1, self.spec_num_tokens) + ) + dsa_plan( + seq_lens_2d=metadata._dsa_seq_lens_2d, + page_size=self.page_size, + out=metadata._dsa_plan, + ) + + def get_cuda_graph_seq_len_fill_value(self): + return self._dense_backend.get_cuda_graph_seq_len_fill_value() + + def advance_draft_forward_metadata(self, seq_lens: torch.Tensor | None = None): + metadata = self.forward_decode_metadata + if metadata is None or metadata.seq_lens_k is None: + raise RuntimeError("DSA draft decode metadata was not initialized") + if seq_lens is None: + metadata.seq_lens_k.add_(1) + else: + metadata.seq_lens_k.copy_(seq_lens[: metadata.seq_lens_k.numel()]) + + dsa_plan( + seq_lens_2d=metadata.seq_lens_k.unsqueeze(1), + page_size=self.page_size, + out=metadata._dsa_plan, + ) + + def init_forward_metadata( + self, + bs: int, + num_extends: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + forward_mode: ForwardMode, + req_to_page: torch.Tensor, + spec_info=None, + **kwargs, + ): + self._dense_backend.init_forward_metadata( + bs=bs, + num_extends=num_extends, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + forward_mode=forward_mode, + req_to_page=req_to_page, + spec_info=spec_info, + **kwargs, + ) + if ( + forward_mode.is_decode() + or forward_mode.is_mixed() + or (forward_mode.is_extend() and self.is_draft) + ): + metadata = self.forward_decode_metadata + # Full-length broadcast: the plan and paged-MQA-logits kernels read only + # the last column, and the per-token causal bound is applied downstream. + metadata._dsa_seq_lens_2d = ( + seq_lens.unsqueeze(1).expand(-1, self.spec_num_tokens).contiguous() + ) + if num_extends < bs: + seq_lens_2d = metadata._dsa_seq_lens_2d[num_extends:] + else: + # The dsa_plan is unused, alias to full-batch seq_lens_2d to generate dsa_plan as a placeholder + seq_lens_2d = metadata._dsa_seq_lens_2d + metadata._dsa_plan = dsa_plan( + seq_lens_2d=seq_lens_2d, page_size=self.page_size + ) + + self._prefill_block_tables = None + if ( + num_extends > 0 + and req_to_page is not None + and forward_mode.is_extend_or_mixed() + ): + cmeta = getattr(self._dense_backend, "chunked_prefill_metadata", None) + cmeta_req_pool_indices = getattr(cmeta, "req_pool_indices", None) + if cmeta is not None and cmeta_req_pool_indices is not None: + ext_idx = cmeta_req_pool_indices[:num_extends].long() + self._prefill_block_tables = req_to_page[ext_idx] + cmeta.block_tables = self._prefill_block_tables + + def _validate_logit_cap(self, logits_soft_cap: float) -> None: + if logits_soft_cap and logits_soft_cap > 0: + raise NotImplementedError( + "TokenSpeed DSA fused dense attention does not support " + f"logits_soft_cap={logits_soft_cap}. Sparse DSA kernels must " + "preserve the capped-score semantics before enabling this model." + ) + + def _validate_dense_context(self, seq_lens: torch.Tensor, bs: int) -> None: + if seq_lens is None or bs <= 0: + return + active_seq_lens = seq_lens[:bs] + if active_seq_lens.numel() == 0: + return + max_seq_len = int(active_seq_lens.max().item()) + if max_seq_len > self.index_topk: + raise NotImplementedError( + "TokenSpeed DSA dense attention is exact only when every " + f"request has seq_len <= index_topk ({self.index_topk}); got " + f"max seq_len {max_seq_len}. Sparse DSA top-k indices are " + "required for longer contexts." + ) + + def _metadata_seq_lens(self, metadata) -> torch.Tensor | None: + seq_lens = getattr(metadata, "seq_lens_k", None) + if seq_lens is not None: + return seq_lens + return getattr(metadata, "seq_lens", None) + + def forward_extend_chunked( + self, + q, + k, + v, + scaling, + logits_soft_cap, + *, + cum_seq_lens_q, + cum_seq_lens_kv, + max_q_len, + max_kv_len, + seq_lens, + batch_size, + causal, + out: torch.Tensor | None = None, + ): + self._validate_logit_cap(logits_soft_cap) + self._validate_dense_context(seq_lens, batch_size) + return self._dense_backend.forward_extend_chunked( + q, + k, + v, + scaling, + logits_soft_cap, + cum_seq_lens_q=cum_seq_lens_q, + cum_seq_lens_kv=cum_seq_lens_kv, + max_q_len=max_q_len, + max_kv_len=max_kv_len, + seq_lens=seq_lens, + batch_size=batch_size, + causal=causal, + out=out, + ) + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer, + out_cache_loc: torch.Tensor, + token_to_kv_pool, + bs: int, + save_kv_cache: bool = True, + topk_indices: torch.Tensor | None = None, + topk_lens: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + self._validate_logit_cap(layer.logit_cap) + if topk_indices is not None: + return self.forward_sparse_decode( + q=q, + k=k, + v=v, + layer=layer, + out_cache_loc=out_cache_loc, + token_to_kv_pool=token_to_kv_pool, + bs=bs, + save_kv_cache=save_kv_cache, + topk_indices=topk_indices, + topk_lens=topk_lens, + ) + metadata = getattr(self, "forward_decode_metadata", None) + seq_lens = self._metadata_seq_lens(metadata) if metadata is not None else None + if seq_lens is not None: + num_extends = int(metadata.num_extends or 0) + self._validate_dense_context(seq_lens[num_extends:], bs) + return self._dense_backend.forward_decode( + q=q, + k=k, + v=v, + layer=layer, + out_cache_loc=out_cache_loc, + token_to_kv_pool=token_to_kv_pool, + bs=bs, + save_kv_cache=save_kv_cache, + **kwargs, + ) + + def forward_sparse_prefill( + self, + *, + q: torch.Tensor, + layer, + token_to_kv_pool, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + workspace_indices: torch.Tensor, + topk_lens: torch.Tensor, + kv_workspace_slots: torch.Tensor | None = None, + max_seq_len: int, + ) -> torch.Tensor: + if layer.logit_cap and layer.logit_cap > 0: + self._validate_logit_cap(layer.logit_cap) + if getattr(token_to_kv_pool, "quant_method", None) == "per_token_head": + raise RuntimeError( + "DSA sparse prefill does not support " + "kv_cache_quant_method='per_token_head' yet." + ) + if workspace_indices.shape[0] != q.shape[0]: + raise RuntimeError( + "DSA sparse prefill metadata token mismatch: " + f"indices={workspace_indices.shape[0]}, q_tokens={q.shape[0]}" + ) + if topk_lens.shape[0] != q.shape[0]: + raise RuntimeError( + "DSA sparse prefill top-k length mismatch: " + f"lens={topk_lens.shape[0]}, q_tokens={q.shape[0]}" + ) + if q.shape[0] == 0: + return q.new_empty((0, layer.tp_q_head_num * layer.v_head_dim)) + if workspace_indices.shape != (q.shape[0], self.index_topk): + raise RuntimeError( + "DSA sparse prefill top-k shape mismatch: " + f"indices={tuple(workspace_indices.shape)}, " + f"expected={(q.shape[0], self.index_topk)}" + ) + if kv_workspace_slots is None: + raise RuntimeError( + "DSA sparse prefill requires kv_workspace_slots to " + "map workspace-local top-k rows back to KV cache slots." + ) + topk_slots = workspace_topk_to_global_slots( + workspace_indices=workspace_indices, + kv_workspace_slots=kv_workspace_slots, + ) + q_view = q.view(q.shape[0], layer.tp_q_head_num, layer.head_dim) + if self.data_type == torch.float8_e4m3fn and q_view.dtype != self.data_type: + q_view = q_view.to(self.data_type) + kv_cache = token_to_kv_pool.get_key_buffer(layer.layer_id) + sparse_kv_cache = None + if hasattr(token_to_kv_pool, "get_sparse_decode_kv_buffer"): + sparse_kv_cache = token_to_kv_pool.get_sparse_decode_kv_buffer( + layer.layer_id + ) + + k_scale = ( + layer.k_scale_float + if getattr(layer, "k_scale_float", None) is not None + else 1.0 + ) + out = dsa_prefill( + q=q_view, + kv_cache=kv_cache, + sparse_kv_cache=sparse_kv_cache, + topk_slots=topk_slots, + topk_lens=topk_lens.to(device=q.device, dtype=torch.int32).contiguous(), + max_seqlen_k=max_seq_len, + 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, + softmax_scale=layer.scaling, + page_size=self.page_size, + logit_cap=layer.logit_cap, + k_scale=k_scale, + ) + return out.reshape(-1, layer.tp_q_head_num * layer.v_head_dim) + + def forward_sparse_decode( + self, + *, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer, + out_cache_loc: torch.Tensor, + token_to_kv_pool, + bs: int, + save_kv_cache: bool, + topk_indices: torch.Tensor, + topk_lens: torch.Tensor | None, + ) -> torch.Tensor: + if self.page_size != 64: + raise RuntimeError( + "DSA sparse decode currently requires page_size=64 for " + f"sparse KV layout, got {self.page_size}." + ) + if getattr(token_to_kv_pool, "quant_method", None) == "per_token_head": + raise RuntimeError( + "DSA sparse decode does not support " + "kv_cache_quant_method='per_token_head' yet." + ) + allow_fp8_query = ( + getattr(self, "data_type", torch.bfloat16) == torch.float8_e4m3fn + and q.dtype == torch.float8_e4m3fn + ) + if q.dtype != torch.bfloat16 and not allow_fp8_query: + raise RuntimeError( + "DSA sparse decode requires BF16 query tensors, or FP8 query " + f"tensors on FP8 KV sparse paths, got {q.dtype}." + ) + if save_kv_cache: + assert k is not None + token_to_kv_pool.set_mla_kv_buffer( + layer, + out_cache_loc, + k[..., : self.kv_lora_rank], + k[..., self.kv_lora_rank :], + ) + + if topk_indices.dtype != torch.int32: + topk_indices = topk_indices.to(torch.int32) + if topk_indices.shape[-1] != self.index_topk: + raise RuntimeError( + "DSA sparse decode top-k width mismatch: " + f"indices={topk_indices.shape[-1]}, expected={self.index_topk}" + ) + num_tokens = q.shape[0] + # Spec-verify feeds q_len_per_req query rows per request while plain + # decode and the draft model's own decode steps feed one; derive the + # width from the actual batch shape (bs is the decode request count) + # rather than spec_num_tokens, which the draft backend inherits from the + # shared config. + if bs > 0 and num_tokens % bs == 0: + q_len_per_req = num_tokens // bs + else: + q_len_per_req = 1 + num_reqs = num_tokens // q_len_per_req + metadata = getattr(self, "forward_decode_metadata", None) + if metadata is None or metadata.seq_lens_k is None: + raise RuntimeError("DSA sparse decode requires decode metadata.") + num_extends = int(metadata.num_extends or 0) + available_reqs = max(0, int(metadata.seq_lens_k.shape[0]) - num_extends) + if available_reqs < num_reqs: + if available_reqs <= 0 or q.shape[0] % available_reqs != 0: + raise RuntimeError( + "DSA sparse decode metadata batch mismatch: " + f"seq_lens={available_reqs}, requests={num_reqs}, " + f"q_tokens={q.shape[0]}." + ) + num_reqs = available_reqs + q_len_per_req = q.shape[0] // available_reqs + seq_lens = metadata.seq_lens_k[num_extends : num_extends + num_reqs] + if seq_lens.numel() != num_reqs: + raise RuntimeError( + "DSA sparse decode metadata batch mismatch: " + f"seq_lens={seq_lens.numel()}, requests={num_reqs}." + ) + num_tokens = q.shape[0] + expected_tokens = num_reqs * int(q_len_per_req) + if num_tokens != expected_tokens: + raise RuntimeError( + "DSA sparse decode token shape mismatch: " + f"q_tokens={num_tokens}, requests={num_reqs}, " + f"q_len_per_req={q_len_per_req}." + ) + if topk_lens is not None: + if topk_lens.dim() != 1 or topk_lens.numel() != num_tokens: + raise RuntimeError( + "DSA sparse decode top-k length mismatch: " + f"lens={tuple(topk_lens.shape)}, q_tokens={num_tokens}." + ) + topk_lens = topk_lens.to(device=q.device, dtype=torch.int32).contiguous() + + q_view = q.view(num_tokens, layer.tp_q_head_num, layer.head_dim) + if self.data_type == torch.float8_e4m3fn: + q_view = q_view.to(self.data_type) + kv_cache = token_to_kv_pool.get_key_buffer(layer.layer_id) + sparse_kv_cache = None + if hasattr(token_to_kv_pool, "get_sparse_decode_kv_buffer"): + sparse_kv_cache = token_to_kv_pool.get_sparse_decode_kv_buffer( + layer.layer_id + ) + + k_scale = ( + layer.k_scale_float + if getattr(layer, "k_scale_float", None) is not None + else 1.0 + ) + max_seqlen_k = int( + getattr(metadata, "max_seq_len_k", 0) or self.max_context_len + ) + out = dsa_decode( + q=q_view, + kv_cache=kv_cache, + sparse_kv_cache=sparse_kv_cache, + topk_slots=topk_indices.view(num_tokens, -1), + topk_lens=topk_lens, + max_seqlen_k=max_seqlen_k, + 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, + softmax_scale=layer.scaling, + page_size=self.page_size, + q_len_per_req=q_len_per_req, + logit_cap=layer.logit_cap, + k_scale=k_scale, + ) + return out.reshape(-1, layer.tp_q_head_num * layer.v_head_dim) + + +register_backend("dsa", {AttentionArch.DSA}, DSABackend) diff --git a/python/tokenspeed/runtime/layers/attention/backends/flashmla.py b/python/tokenspeed/runtime/layers/attention/backends/flashmla.py new file mode 100644 index 0000000..a708ff4 --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/backends/flashmla.py @@ -0,0 +1,852 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch +from tokenspeed_kernel.ops.attention.flash_attn import flash_attn_varlen_func +from tokenspeed_kernel.ops.attention.flash_mla import ( + flash_mla_with_kvcache, + get_mla_metadata, +) +from tokenspeed_kernel.ops.attention.flashinfer import ( + BatchMLAPagedAttentionWrapper, + BatchPrefillWithRaggedKVCacheWrapper, +) + +from tokenspeed.runtime.configs.model_config import AttentionArch +from tokenspeed.runtime.execution.forward_batch_info import ForwardMode +from tokenspeed.runtime.layers.attention.backends.base import AttentionBackend +from tokenspeed.runtime.layers.attention.chunk import ( + build_chunked_prefill_metadata_arrays, +) +from tokenspeed.runtime.layers.attention.configs.mla import MLAConfig +from tokenspeed.runtime.layers.attention.registry import register_backend +from tokenspeed.runtime.layers.attention.utils import ( + create_flashinfer_kv_indices_triton, +) +from tokenspeed.runtime.spec_decode.eagle import ( + EagleDraftInput, + generate_attn_arg_prefill, +) +from tokenspeed.runtime.utils.env import global_server_args_dict +from tokenspeed.runtime.utils.flashinfer_config import get_flashinfer_workspace_size + +PAGE_SIZE = 64 + +if TYPE_CHECKING: + from tokenspeed.runtime.layers.paged_attention import PagedAttention + + +@dataclass +class FlashMLADecodeMetadata: + num_extends: int = 0 + flashmla_metadata: tuple | None = None + num_splits: torch.Tensor | None = None + block_table: torch.Tensor | None = None + + +@dataclass +class _PrefillMetadata: + prefill_wrapper: BatchMLAPagedAttentionWrapper + use_ragged: bool + + +@dataclass +class _ChunkedPrefillMetadata: + extend_prefix_lens: torch.Tensor + extend_prefix_lens_cpu: torch.Tensor + extend_seq_lens: torch.Tensor + extend_seq_lens_cpu: torch.Tensor + req_pool_indices: torch.Tensor + cum_extend_seq_lens: torch.Tensor + max_extend_seq_len: int + chunked_loop_num: int + chunk_kv_indices_list: list + chunked_seq_len: torch.Tensor + cu_chunked_seq_len: torch.Tensor + max_chunk_len_per_loop: list + + +# Shared across all flashinfer prefill wrappers used by FlashMLABackend. +_global_workspace_buffer = None + + +class FlashMLABackend(AttentionBackend): + """FlashMLA attention backend for TokenSpeed scheduling. + + Uses the FlashMLA kernel for decode (any q_len); uses FlashInfer's MLA + prefill wrappers for the EXTEND path. + """ + + def __init__(self, config: MLAConfig): + super().__init__(config) + + # Parse constants + self.max_context_len = config.context_len + self.kv_cache_quant_method = config.kv_cache_quant_method + self.cache_dtype = config.kv_cache_dtype + + # 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 = config.kv_lora_rank + config.qk_rope_head_dim + self.scaling = config.scaling + self.softmax_scale = config.scaling + self.data_type = config.kv_cache_dtype + self.q_data_type = config.dtype + self.num_local_heads = config.num_attention_heads // config.attn_tp_size + self.num_q_heads = config.num_attention_heads // config.attn_tp_size + + # FlashMLA-specific + self.draft_token_num = 0 + + if self.kv_cache_quant_method == "per_token_head": + raise NotImplementedError( + "FlashMLABackend no longer supports " + "kv_cache_quant_method='per_token_head'." + ) + if self.cache_dtype == torch.float8_e4m3fn: + raise NotImplementedError( + "FlashMLABackend no longer supports dense FP8 KV cache. " + "Use a non-FP8 KV cache." + ) + + # Workspace buffer + flashinfer prefill wrappers (EXTEND path only). + global _global_workspace_buffer + if _global_workspace_buffer is None: + _global_workspace_buffer = torch.empty( + get_flashinfer_workspace_size(), + dtype=torch.uint8, + device=config.device, + ) + self.workspace_buffer = _global_workspace_buffer + + max_bs = config.max_bs + self.kv_indptr = torch.zeros( + (max_bs + 1,), dtype=torch.int32, device=config.device + ) + self.qo_indptr = torch.zeros( + (max_bs + 1,), dtype=torch.int32, device=config.device + ) + + self.prefill_wrapper_ragged = BatchPrefillWithRaggedKVCacheWrapper( + self.workspace_buffer, "NHD" + ) + self.prefill_wrapper_paged = BatchMLAPagedAttentionWrapper( + self.workspace_buffer, + backend="auto", + ) + self.indices_updater_prefill = _PrefillIndicesUpdater(config, self) + + # Metadata state. Decode and prefill metadata are split so MIXED batches + # can carry both simultaneously (decode-half + prefill-half sub-contexts + # dispatch to their respective metadata). + self.forward_decode_metadata: FlashMLADecodeMetadata | None = None + self.forward_prefill_metadata: _PrefillMetadata | None = None + self.chunked_prefill_metadata: _ChunkedPrefillMetadata | None = None + self.last_seq_lens_sum: int | None = None + + # ------------------------------------------------------------------ + # Metadata init + # ------------------------------------------------------------------ + + def init_forward_metadata( + self, + bs: int, + num_extends: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + forward_mode: ForwardMode, + req_to_page: torch.Tensor = None, + extend_with_prefix: bool = False, + extend_prefix_lens: torch.Tensor | None = None, + spec_info=None, + **kwargs, + ): + if forward_mode.is_extend_or_mixed(): + self._init_prefill_metadata( + req_pool_indices=req_pool_indices[:num_extends], + seq_lens=seq_lens[:num_extends], + req_to_page=req_to_page, + extend_with_prefix=extend_with_prefix, + extend_prefix_lens=extend_prefix_lens, + extend_prefix_lens_cpu=kwargs.pop("extend_prefix_lens_cpu"), + extend_seq_lens=kwargs.pop("extend_seq_lens"), + extend_seq_lens_cpu=kwargs.pop("extend_seq_lens_cpu"), + ) + # Under is_draft, also fill decode_metadata under any forward_mode so + # the drafter's multi-step loop has metadata. Wrapper pre-writes + # draft_seq_lens before calling here, so `seq_lens` aliases the + # drafter's live buffer for step-1+ advances. + if ( + forward_mode.is_decode_or_idle() + or forward_mode.is_mixed() + or (forward_mode.is_extend() and self.is_draft) + ): + self._init_decode_metadata( + bs, num_extends, req_pool_indices, seq_lens, req_to_page + ) + + @contextmanager + def override_num_extends(self, num_extends: int): + assert self.forward_decode_metadata is not None + prev = self.forward_decode_metadata.num_extends + self.forward_decode_metadata.num_extends = num_extends + try: + yield + finally: + self.forward_decode_metadata.num_extends = prev + + def _init_decode_metadata( + self, + bs: int, + num_extends: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + req_to_page: torch.Tensor, + ): + if req_to_page is not None: + block_table = req_to_page[req_pool_indices] + else: + block_table = None + + # When spec-dec is active (self.spec_num_tokens > 1), advance per-row + # seq_lens by the worst-case verify width so the tile planner covers + # the longest path. + if self.spec_num_tokens > 1: + plan_seq_lens = seq_lens + self.draft_token_num + num_heads_plan = self.draft_token_num * self.num_q_heads + else: + plan_seq_lens = seq_lens + num_heads_plan = self.num_q_heads + + mla_metadata, num_splits = get_mla_metadata( + plan_seq_lens.to(torch.int32), + num_heads_plan, + 1, + ) + self.forward_decode_metadata = FlashMLADecodeMetadata( + num_extends=num_extends, + flashmla_metadata=mla_metadata, + num_splits=num_splits, + block_table=block_table, + ) + + def _init_prefill_metadata( + self, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + req_to_page: torch.Tensor, + extend_with_prefix: bool, + extend_prefix_lens: torch.Tensor | None, + extend_prefix_lens_cpu: torch.Tensor, + extend_seq_lens: torch.Tensor, + extend_seq_lens_cpu: torch.Tensor, + ): + # EXTEND path — flashinfer ragged/paged prefill. + if extend_prefix_lens is None: + raise RuntimeError( + "FlashMLABackend.init_forward_metadata requires " + "extend_prefix_lens in extend mode." + ) + seq_lens_cpu = seq_lens.cpu() + seq_lens_sum = seq_lens_cpu.sum().item() + self.last_seq_lens_sum = seq_lens_sum + + extend_no_prefix = not extend_with_prefix + use_ragged = ( + not global_server_args_dict["mla_disable_ragged"] and extend_no_prefix + ) + + self.indices_updater_prefill.update( + req_pool_indices, + seq_lens, + seq_lens_sum, + extend_prefix_lens, + req_to_page=req_to_page, + prefill_wrapper_paged=self.prefill_wrapper_paged, + use_ragged=use_ragged, + ) + self.forward_prefill_metadata = _PrefillMetadata( + self.prefill_wrapper_paged, use_ragged + ) + + num_extends = extend_seq_lens.shape[0] + cum_extend_seq_lens = torch.zeros( + num_extends + 1, device=self.device, dtype=torch.int32 + ) + torch.cumsum(extend_seq_lens, dim=0, out=cum_extend_seq_lens[1:]) + max_extend_seq_len = extend_seq_lens_cpu.max().item() + ( + chunked_loop_num, + chunk_kv_indices_list, + chunked_seq_len, + cu_chunked_seq_len, + max_chunk_len_per_loop, + ) = build_chunked_prefill_metadata_arrays( + extend_prefix_lens, + extend_prefix_lens_cpu, + req_to_page, + req_pool_indices, + PAGE_SIZE, + ) + self.chunked_prefill_metadata = _ChunkedPrefillMetadata( + extend_prefix_lens=extend_prefix_lens, + extend_prefix_lens_cpu=extend_prefix_lens_cpu, + extend_seq_lens=extend_seq_lens, + extend_seq_lens_cpu=extend_seq_lens_cpu, + req_pool_indices=req_pool_indices, + cum_extend_seq_lens=cum_extend_seq_lens, + max_extend_seq_len=max_extend_seq_len, + chunked_loop_num=chunked_loop_num, + chunk_kv_indices_list=chunk_kv_indices_list, + chunked_seq_len=chunked_seq_len, + cu_chunked_seq_len=cu_chunked_seq_len, + max_chunk_len_per_loop=max_chunk_len_per_loop, + ) + + # ------------------------------------------------------------------ + # CUDA graph (decode only, any q_len) + # ------------------------------------------------------------------ + + def init_cuda_graph_state(self, max_bs: int, seq_lens_buf: torch.Tensor): + del seq_lens_buf # flashmla allocates its own buffers. + max_context_len = self.max_context_len + PAGE_SIZE - 1 + # 4 PAGES are reserved for speculation + cuda_graph_kv_indices = torch.full( + (max_bs, (max_context_len + 4 * PAGE_SIZE) // PAGE_SIZE), + 1, + dtype=torch.int32, + device="cuda", + ) + + if self.draft_token_num: + ( + 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.draft_token_num * 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, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + forward_mode: ForwardMode, + ): + block_table = self.cuda_graph_kv_indices[:bs] + is_target_verify = ( + forward_mode.is_decode_or_idle() + and not self.is_draft + and self.spec_num_tokens > 1 + ) + is_draft_extend = ( + forward_mode.is_decode_or_idle() + and self.is_draft + and self.spec_num_tokens > 1 + ) + + if forward_mode.is_decode_or_idle() and self.spec_num_tokens == 1: + 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.cuda_graph_kv_indices[:bs].copy_(block_table) + self.forward_decode_metadata = FlashMLADecodeMetadata( + num_extends=0, + flashmla_metadata=self.cuda_graph_mla_metadata, + num_splits=self.cuda_graph_num_splits[: bs + 1], + block_table=self.cuda_graph_kv_indices[:bs, :], + ) + elif is_target_verify or is_draft_extend: + seq_lens = seq_lens + self.draft_token_num + mla_metadata, num_splits = get_mla_metadata( + seq_lens.to(torch.int32), + self.draft_token_num * self.num_q_heads, + 1, + ) + self.cuda_graph_mla_metadata.copy_(mla_metadata) + self.cuda_graph_num_splits[: bs + 1].copy_(num_splits) + self.cuda_graph_kv_indices[:bs].copy_(block_table) + self.forward_decode_metadata = FlashMLADecodeMetadata( + num_extends=0, + flashmla_metadata=self.cuda_graph_mla_metadata, + num_splits=self.cuda_graph_num_splits[: bs + 1], + block_table=self.cuda_graph_kv_indices[:bs], + ) + else: + raise RuntimeError(f"Not supported forward mode: {forward_mode}") + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + forward_mode: ForwardMode = None, + req_to_page: torch.Tensor = None, + **kwargs, + ): + if forward_mode is None or not forward_mode.is_decode_or_idle(): + raise RuntimeError(f"Not supported forward mode: {forward_mode}") + + req_pool_indices = req_pool_indices[:bs] + if req_to_page is not None: + block_table = req_to_page[req_pool_indices] + else: + block_table = self.cuda_graph_kv_indices[:bs] + seq_lens = seq_lens[:bs] + + is_target_verify = not self.is_draft and self.spec_num_tokens > 1 + is_draft_extend = self.is_draft and self.spec_num_tokens > 1 + + if self.spec_num_tokens == 1: + mla_metadata, num_splits = get_mla_metadata( + seq_lens.to(torch.int32), + self.num_q_heads, + 1, + ) + elif is_target_verify or is_draft_extend: + seq_lens = seq_lens + self.draft_token_num + mla_metadata, num_splits = get_mla_metadata( + seq_lens.to(torch.int32), + self.draft_token_num * self.num_q_heads, + 1, + ) + else: + raise RuntimeError(f"Not supported forward mode: {forward_mode}") + + self.cuda_graph_mla_metadata.copy_(mla_metadata) + self.cuda_graph_num_splits[: bs + 1].copy_(num_splits) + self.cuda_graph_kv_indices[:bs].copy_(block_table) + self.forward_decode_metadata.num_extends = 0 + self.forward_decode_metadata.flashmla_metadata = self.cuda_graph_mla_metadata + self.forward_decode_metadata.num_splits = self.cuda_graph_num_splits[: bs + 1] + self.forward_decode_metadata.block_table = self.cuda_graph_kv_indices[:bs] + + def get_cuda_graph_seq_len_fill_value(self): + return 1 + + # ------------------------------------------------------------------ + # Forward + # ------------------------------------------------------------------ + + def forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: PagedAttention, + out_cache_loc: torch.Tensor, + token_to_kv_pool, + bs: int, + save_kv_cache: bool = True, + seq_lens: torch.Tensor | None = None, + forward_mode: ForwardMode | None = None, + **kwargs, + ): + q_len_per_req = q.shape[0] // bs if bs > 0 else 1 + is_target_verify = ( + forward_mode is not None + and forward_mode.is_decode_or_idle() + and not self.is_draft + and q_len_per_req > 1 + ) + is_draft_extend = ( + forward_mode is not None + and forward_mode.is_decode_or_idle() + and self.is_draft + and q_len_per_req > 1 + ) + + if forward_mode is None or forward_mode.is_extend(): + # Prefill: dispatch to ragged (MHA-style) or absorbed (MQA) path. + if self.forward_prefill_metadata.use_ragged: + return self._forward_normal_extend(q, k, v, layer, save_kv_cache) + else: + return self._forward_absorbed_extend( + q, + k, + v, + layer, + out_cache_loc, + token_to_kv_pool, + save_kv_cache, + ) + + assert is_target_verify or is_draft_extend + if k is not None: + assert v is not None + if save_kv_cache: + token_to_kv_pool.set_kv_buffer(layer, out_cache_loc, k, v) + + metadata = self.forward_decode_metadata + num_extends = metadata.num_extends + bs = ( + q.shape[0] + if is_draft_extend + else metadata.block_table.shape[0] - num_extends + ) + k_cache = token_to_kv_pool.get_key_buffer(layer.layer_id) + + assert ( + layer.tp_q_head_num == self.num_q_heads + ), f"{layer.tp_q_head_num=} != {self.num_q_heads=}" + reshape_q = q.view(bs, -1, self.num_q_heads, layer.head_dim) + + o, _ = flash_mla_with_kvcache( + q=reshape_q, + k_cache=k_cache.view(-1, PAGE_SIZE, 1, self.kv_cache_dim), + block_table=metadata.block_table[num_extends : num_extends + bs], + cache_seqlens=seq_lens.to(torch.int32) + self.draft_token_num, + head_dim_v=self.kv_lora_rank, + tile_scheduler_metadata=metadata.flashmla_metadata, + num_splits=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_chunked( + self, + q, + k, + v, + scaling, + logits_soft_cap=None, + *, + cum_seq_lens_q, + cum_seq_lens_kv, + max_q_len, + max_kv_len, + seq_lens, + batch_size, + causal, + out: torch.Tensor | None = None, + ): + if causal: + step_counter = getattr(self, "step_counter", None) + if step_counter is not None: + step_counter.record_cache() + head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + # flash_attn_varlen_func has no `out=` parameter; copy into the + # caller-provided buffer at the end when requested. + output, lse, *_ = flash_attn_varlen_func( + q=q.view(-1, self.num_local_heads, head_dim), + k=k.view(-1, self.num_local_heads, head_dim).to(q.dtype), + v=v.view(-1, self.num_local_heads, self.v_head_dim).to(q.dtype), + cu_seqlens_q=cum_seq_lens_q, + cu_seqlens_k=cum_seq_lens_kv, + max_seqlen_q=max_q_len, + max_seqlen_k=max_kv_len, + softmax_scale=scaling, + causal=causal, + return_attn_probs=True, + ) + if out is not None: + out.copy_(output.view(out.shape)) + output = out + # lse must be transposed when using fa3. + return output, lse.T.contiguous() + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: PagedAttention, + out_cache_loc: torch.Tensor, + token_to_kv_pool, + bs: int, + save_kv_cache: bool = True, + seq_lens: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + # Multi-token decode (target verify or drafter compound) reuses + # the multi-token kernel path in forward_extend. + q_len_per_req = q.shape[0] // bs if bs > 0 else 1 + if q_len_per_req > 1: + return self.forward_extend( + q, + k, + v, + layer, + out_cache_loc, + token_to_kv_pool, + bs, + save_kv_cache=save_kv_cache, + seq_lens=seq_lens, + forward_mode=ForwardMode.DECODE, + **kwargs, + ) + + if k is not None: + assert v is not None + if save_kv_cache: + token_to_kv_pool.set_kv_buffer( + layer, + out_cache_loc, + k, + v, + ) + bs = q.shape[0] + metadata = self.forward_decode_metadata + num_extends = metadata.num_extends + k_cache = token_to_kv_pool.get_key_buffer(layer.layer_id) + assert ( + layer.tp_q_head_num == self.num_q_heads + ), f"{layer.tp_q_head_num=} != {self.num_q_heads=}" + reshape_q = q.view(bs, -1, self.num_q_heads, layer.head_dim) + cache_lens = seq_lens + + o, _ = flash_mla_with_kvcache( + q=reshape_q, + k_cache=k_cache.view(-1, PAGE_SIZE, 1, self.kv_cache_dim), + block_table=metadata.block_table[num_extends : num_extends + bs], + cache_seqlens=cache_lens.to(torch.int32), + head_dim_v=self.kv_lora_rank, + tile_scheduler_metadata=metadata.flashmla_metadata, + num_splits=metadata.num_splits, + softmax_scale=layer.scaling, + causal=True, + ) + + return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) + + # ------------------------------------------------------------------ + # EXTEND prefill helpers + # ------------------------------------------------------------------ + + def _forward_normal_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: PagedAttention, + save_kv_cache: bool = True, + ): + assert not save_kv_cache + + o = self.prefill_wrapper_ragged.forward( + q, + k.view(-1, layer.tp_k_head_num, layer.head_dim), + v.view(-1, layer.tp_k_head_num, layer.v_head_dim), + causal=True, + sm_scale=layer.scaling, + logits_soft_cap=layer.logit_cap, + ) + return o.view(-1, layer.tp_q_head_num * layer.v_head_dim) + + def _forward_absorbed_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: PagedAttention, + out_cache_loc: torch.Tensor, + token_to_kv_pool, + save_kv_cache: bool = True, + ): + # q is whole Q [T, H, head_dim]; k is whole latent [T, 1, head_dim]. + # flashinfer prefill_wrapper.run() requires q_nope / q_pe split, so + # slice views here (free) before handing off to the kernel. + assert k is not None + + if save_kv_cache: + token_to_kv_pool.set_mla_kv_buffer( + layer, + out_cache_loc, + k[..., : layer.v_head_dim], + k[..., layer.v_head_dim :], + ) + + q = q.view(-1, layer.tp_q_head_num, layer.head_dim) + q_nope = q[..., : layer.v_head_dim] + q_pe = q[..., layer.v_head_dim :] + o = q_nope.new_empty(q_nope.shape) + + k_buf = token_to_kv_pool.get_key_buffer(layer.layer_id).to(q_nope.dtype) + o = self.forward_prefill_metadata.prefill_wrapper.run( + q_nope, + q_pe, + 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) + + +class _PrefillIndicesUpdater: + """Plans FlashInfer MLA prefill wrappers for the EXTEND path.""" + + def __init__(self, config: MLAConfig, attn_backend: FlashMLABackend): + self.num_local_heads = config.num_attention_heads // config.attn_tp_size + self.kv_cache_quant_method = config.kv_cache_quant_method + 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.scaling = config.scaling + self.data_type = config.kv_cache_dtype + self.q_data_type = config.dtype + self.attn_backend = attn_backend + + self.kv_indptr = attn_backend.kv_indptr + self.qo_indptr = attn_backend.qo_indptr + self.prefill_wrapper_ragged = attn_backend.prefill_wrapper_ragged + + def update( + self, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + seq_lens_sum: int, + prefix_lens: torch.Tensor, + req_to_page: torch.Tensor = None, + prefill_wrapper_paged: BatchMLAPagedAttentionWrapper = None, + use_ragged: bool = False, + spec_info: EagleDraftInput | None = None, + ): + if use_ragged: + paged_kernel_lens = prefix_lens + paged_kernel_lens_sum = 0 + 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, + req_to_page=req_to_page, + spec_info=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, + req_to_page: torch.Tensor = None, + spec_info: EagleDraftInput | None = None, + ): + bs = len(seq_lens) + sm_scale = self.scaling + + if spec_info is None: + assert len(seq_lens) == len(req_pool_indices) + torch.cumsum(paged_kernel_lens, dim=0, out=kv_indptr[1 : bs + 1]) + kv_indptr = kv_indptr[: bs + 1] + if wrapper_paged._use_cuda_graph: + kv_indices = wrapper_paged._kv_indices_buf + else: + kv_indices = torch.empty( + paged_kernel_lens_sum, + dtype=torch.int32, + device=req_pool_indices.device, + ) + if req_to_page is not None: + create_flashinfer_kv_indices_triton[(bs,)]( + req_to_page, + req_pool_indices, + paged_kernel_lens, + kv_indptr, + None, + kv_indices, + req_to_page.shape[1], + ) + torch.cumsum(seq_lens - prefix_lens, dim=0, out=qo_indptr[1 : bs + 1]) + qo_indptr = qo_indptr[: bs + 1] + else: + kv_indices, kv_indptr, qo_indptr, _ = generate_attn_arg_prefill( + spec_info.draft_token_num, + req_pool_indices, + paged_kernel_lens, + req_to_page, + ) + + if use_ragged: + 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, + ) + else: + 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, + ) + + +register_backend("flashmla", {AttentionArch.MLA}, FlashMLABackend) diff --git a/python/tokenspeed/runtime/layers/attention/backends/flat_groups.py b/python/tokenspeed/runtime/layers/attention/backends/flat_groups.py new file mode 100644 index 0000000..5f5817a --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/backends/flat_groups.py @@ -0,0 +1,331 @@ +"""Shared flat KV-cache group machinery for attention backends. + +A flat-capable backend (``uses_flat_cache_groups = True``) receives one page +table per cache group (``flat_block_tables: dict[group_id, [bs, max_pages]]``) +instead of the radix single table, and must route every KV read AND write +through the layer's own group (M-W1). This mixin holds the group-selection, +write-location, and CUDA-graph per-group buffer machinery shared by the MHA +and TRT-LLM backends; model/kernel-specific constraints (spec decode, DFLASH) +stay in the backends. + +Table contract (canonical): rows are requests (padded rows carry the +zero-init dummy page 0), column tails pad with -1 and are never read past +``cache_seqlens``; SWA holes sit only at the window front and are written as +the null page 0 by the scheduler export. +""" + +from __future__ import annotations + +import os + +import torch + + +class FlatCacheGroupsMixin: + """Per-group table/write-loc selection + CUDA-graph buffer discipline. + + Host class requirements: ``self.device``, ``self.page_size``, + ``self.max_num_pages``, ``self.forward_decode_metadata`` (with + ``page_tables``/``out_cache_locs`` fields), and calling + :meth:`_init_flat_graph_buffers` from ``init_cuda_graph_state``. + """ + + # family="state" group ids (GDN/mamba state pages); learned from the + # pool's specs in init_cuda_graph_state, shed from every table here. + flat_state_group_ids: frozenset[str] = frozenset() + + # Value for CUDA-graph buffer column tails past this replay's table + # width. -1 is a debug tripwire (never read past cache_seqlens by the + # MHA kernels); backends whose kernels assume a full-width table + # (trtllm: row stride derived from max_kv_len) override with 0, the + # zero-init dummy page — always safe to dereference. + flat_table_tail_pad: int = -1 + + # ------------------------------------------------------------------ + # Group selection + # ------------------------------------------------------------------ + + @staticmethod + def _select_group_entry(layer, mapping, what: str): + """Pick this layer's entry from a flat per-group dict (page tables or + write locs): the layer's group entry, or the sole entry when the + layer carries no/unknown group id. TODO(radix-removal): collapses to + `mapping[layer.group_id]` once flat is the only path. + """ + group_id = getattr(layer, "group_id", "") + if not group_id or group_id not in mapping: + if len(mapping) == 1: + return next(iter(mapping.values())) + raise KeyError( + f"{what}: layer group_id={group_id!r} not in flat group " + f"keys {sorted(mapping)}" + ) + return mapping[group_id] + + def _select_page_table(self, layer, metadata): + if metadata.page_tables is None: + return metadata.page_table + return self._select_group_entry(layer, metadata.page_tables, "page table") + + def _select_out_cache_loc( + self, layer, metadata, out_cache_loc, prefer_caller=False + ): + # prefer_caller: draft chains own per-step locs; metadata's single loc would pin every step to one slot. + if metadata.out_cache_locs is None or prefer_caller: + return out_cache_loc + return self._select_group_entry( + layer, metadata.out_cache_locs, "flat write locs" + ) + + def _prewrite_metadata(self, forward_mode): + """Metadata slot the fused prewrite writes against. Default: the + decode slot (MHA gates prewrite to decode); backends that prewrite + on extend too (trtllm) override to pick their extend/prefill slot. + """ + return self.forward_decode_metadata + + def select_out_cache_loc(self, layer, out_cache_loc, forward_mode=None): + """Per-group write locations for out-of-backend KV writers (fused + RoPE prewrite): the write must land in the pages this layer's group + reads, never the scheduler's single-table locations. + """ + metadata = self._prewrite_metadata(forward_mode) + if metadata is None or metadata.out_cache_locs is None: + return out_cache_loc + return self._select_out_cache_loc(layer, metadata, out_cache_loc) + + def _shed_state_groups(self, tables): + """Drop family="state" groups (GDN/mamba state pages, consumed by the + mamba backend): computing write locs / capture buffers over the + hole-heavy state table writes the dummy page and trips + TOKENSPEED_FLAT_DEBUG. Returns None when nothing is left. + """ + if not tables: + return None + if self.flat_state_group_ids: + tables = { + gid: table + for gid, table in tables.items() + if gid not in self.flat_state_group_ids + } + return tables or None + + def _learn_flat_state_groups(self, paged_cache_group_specs) -> None: + """Record the pool's family="state" group ids (see + flat_state_group_ids); called from init_cuda_graph_state, the one + place the pool's specs reach every backend.""" + self.flat_state_group_ids = frozenset( + str(spec.group_id) + for spec in paged_cache_group_specs + if spec.family == "state" + ) + + # ------------------------------------------------------------------ + # Write locations + # ------------------------------------------------------------------ + + @staticmethod + def _compute_flat_decode_out_cache_locs( + page_tables, seq_lens, page_size, num_tokens_per_req=1 + ): + """Per-group decode write locs, gathered from the group's own read + table (M-W1). Plain decode writes one token per request at seq_len-1; + spec verify writes num_tokens_per_req at seq_len-N..seq_len-1, + flattened token-major per request ([bs*N], radix verify layout). + Positions clamp at 0 for graph-padded rows (seq_len 1 < N), which + dereference the dummy page harmlessly. The tail page is never a hole + (SWA holes sit only at the window front). + """ + n = num_tokens_per_req + if n == 1: + pos = (seq_lens - 1).to(torch.int64) + else: + steps = torch.arange(n, device=seq_lens.device, dtype=torch.int64) + pos = (seq_lens.to(torch.int64).unsqueeze(1) - n + steps).clamp_min(0) + pos = pos.reshape(-1) + page_idx = pos // page_size + off = (pos % page_size).to(torch.int32) + out = {} + for gid, table in page_tables.items(): + if n == 1: + pages = table.gather(1, page_idx.unsqueeze(1)).squeeze(1) + else: + pages = table.gather(1, page_idx.view(-1, n)).reshape(-1) + out[gid] = pages * page_size + off + return out + + @staticmethod + def _compute_flat_extend_out_cache_locs( + page_tables, extend_prefix_lens_cpu, extend_seq_lens_cpu, page_size + ): + """Per-group extend write locs: positions [prefix_len, seq_len) per + request, flattened in q/k/v token order (cu_extend_seq_lens). Bounds + come from the CPU mirrors — no per-request GPU sync. + TODO(flat-perf): batch the per-request loop via repeat_interleave. + """ + device = next(iter(page_tables.values())).device + prefix_lens = [int(x) for x in extend_prefix_lens_cpu.tolist()] + extend_lens = [int(x) for x in extend_seq_lens_cpu.tolist()] + out = {gid: [] for gid in page_tables} + for i, (start, num_new) in enumerate(zip(prefix_lens, extend_lens)): + pos = torch.arange(start, start + num_new, dtype=torch.int64, device=device) + page_idx = pos // page_size + off = (pos % page_size).to(torch.int32) + for gid, table in page_tables.items(): + pages = table[i].gather(0, page_idx) + out[gid].append(pages * page_size + off) + return { + gid: ( + torch.cat(chunks) + if chunks + else torch.empty(0, dtype=torch.int32, device=device) + ) + for gid, chunks in out.items() + } + + @staticmethod + def _maybe_check_flat_write_locs(page_tables, out_cache_locs, page_size): + """TOKENSPEED_FLAT_DEBUG=1 (eager only, GPU sync): write pages must + be real and inside the group's table. Not for graph-padded batches — + dummy rows would trip the non-hole assert (see the padding contract + in _flat_replay_fill). + """ + if os.environ.get("TOKENSPEED_FLAT_DEBUG") != "1": + return + for gid, locs in out_cache_locs.items(): + pages = (locs // page_size).to(torch.int32) + table = page_tables[gid] + assert ( + pages != 0 + ).all(), f"flat write loc in null page 0 for group {gid!r}" + real = table[table > 0] + assert torch.isin( + pages, real + ).all(), f"flat write pages escape group {gid!r}'s table" + + # ------------------------------------------------------------------ + # CUDA-graph per-group buffers + # ------------------------------------------------------------------ + + def _init_flat_graph_buffers(self, max_bs: int) -> None: + """Reset the lazily-allocated per-group persistent buffers; call from + init_cuda_graph_state BEFORE any backend early return — replay reads + the dict unconditionally for the stale-table guard.""" + self.cuda_graph_flat_page_tables: dict[str, torch.Tensor] = {} + self.cuda_graph_flat_out_cache_locs: dict[str, torch.Tensor] = {} + self._cuda_graph_max_bs = max_bs + + def _flat_capture_group_views( + self, bs: int, flat_cache_group_ids, tokens_per_req: int = 1 + ): + """Capture-time (page_tables, out_cache_locs) per-group views into the + persistent buffers, lazily allocated. Real tables only arrive at + replay, which copy_s fresh data to these graph-recorded addresses. + Verify (tokens_per_req = spec_num_tokens) keeps [bs]-row tables but + records [bs*N] write-loc views (token-major, radix verify layout). + Returns (None, None) when only state groups (or none) are delivered. + """ + if not flat_cache_group_ids: + return None, None + page_tables = {} + out_cache_locs = {} + for gid in flat_cache_group_ids: + if gid in self.flat_state_group_ids: + # State pages ride to the mamba backend; no buffers here. + continue + buf = self.cuda_graph_flat_page_tables.get(gid) + if buf is None: + buf = torch.zeros( + (self._cuda_graph_max_bs, self.max_num_pages), + dtype=torch.int32, + device=self.device, + ) + self.cuda_graph_flat_page_tables[gid] = buf + loc_buf = self.cuda_graph_flat_out_cache_locs.get(gid) + if ( + loc_buf is None + or loc_buf.shape[0] < self._cuda_graph_max_bs * tokens_per_req + ): + loc_buf = torch.zeros( + (self._cuda_graph_max_bs * tokens_per_req,), + dtype=torch.int32, + device=self.device, + ) + self.cuda_graph_flat_out_cache_locs[gid] = loc_buf + page_tables[gid] = buf[:bs, :] + out_cache_locs[gid] = loc_buf[: bs * tokens_per_req] + if not page_tables: + # Only state groups delivered: nothing for this backend. + return None, None + return page_tables, out_cache_locs + + def _flat_replay_stale_guard(self, bs: int, flat_block_tables) -> None: + """Fail loudly instead of replaying over stale/zero page tables. + bs == 0 may skip: col-0 buffer entries stay valid (never -1), + outputs are discarded, and only unit tests reach it.""" + if not self.cuda_graph_flat_page_tables or bs <= 0: + return + name = type(self).__name__ + if not flat_block_tables: + raise RuntimeError( + f"{name} replay: flat per-group CUDA-graph buffers " + f"exist for groups " + f"{sorted(self.cuda_graph_flat_page_tables)} " + f"but flat_block_tables is missing/empty at bs={bs}; the " + "captured graph would read stale page tables." + ) + missing = set(self.cuda_graph_flat_page_tables) - set(flat_block_tables) + if missing: + raise RuntimeError( + f"{name} replay: flat_block_tables at bs=" + f"{bs} is missing captured groups {sorted(missing)} " + f"(delivered: {sorted(flat_block_tables)}); the captured " + "graph would read stale page tables for those groups." + ) + + def _flat_replay_fill( + self, bs: int, flat_block_tables, seq_lens, tokens_per_req: int = 1 + ) -> None: + """Copy this replay's tables into the captured buffers and recompute + the per-group write locs from the live seq_lens (tokens_per_req locs + per request on the spec-verify path). + + Padding contract (canonical; bs is the padded bs): dummy ROWS pad + with 0 — replayed at seq_lens=1 they dereference exactly col 0, + the zero-init dummy page. Column tails pad with -1, never read + past cache_seqlens. + """ + for gid, src in flat_block_tables.items(): + if gid in self.flat_state_group_ids: + # State group: the mamba backend consumes it directly. + continue + buf = self.cuda_graph_flat_page_tables[gid] + cols = src.shape[1] + # cols >= 1: a zero-width table would leave dummy rows' + # col 0 unwritten. + assert 1 <= cols <= buf.shape[1], ( + f"flat table for group {gid!r}: {cols} cols outside" + f" [1, {buf.shape[1]}] (CUDA-graph buffer width)" + ) + assert src.shape[0] >= bs, ( + f"flat table for group {gid!r} has {src.shape[0]} rows" + f" < padded bs {bs}" + ) + buf[:bs, :cols].copy_(src[:bs, :]) + if cols < buf.shape[1]: + buf[:bs, cols:].fill_(self.flat_table_tail_pad) + + # seq_lens is the controller-filled live buffer (current lens + + # padding 1s), written BEFORE replay init, so [:bs] is current. + locs = self._compute_flat_decode_out_cache_locs( + { + gid: self.cuda_graph_flat_page_tables[gid][:bs, :] + for gid in flat_block_tables + if gid not in self.flat_state_group_ids + }, + seq_lens[:bs], + self.page_size, + tokens_per_req, + ) + for gid, val in locs.items(): + self.cuda_graph_flat_out_cache_locs[gid][: bs * tokens_per_req].copy_(val) diff --git a/python/tokenspeed/runtime/layers/attention/backends/hybrid_linear_attn.py b/python/tokenspeed/runtime/layers/attention/backends/hybrid_linear_attn.py new file mode 100644 index 0000000..728c668 --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/backends/hybrid_linear_attn.py @@ -0,0 +1,1691 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Hybrid linear attention backend for Qwen3.5 GDN models.""" + +from __future__ import annotations + +import os +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch +from tokenspeed_kernel.ops.attention import GdnCheckpointLayout, gdn_chunk_prefill +from tokenspeed_kernel.ops.attention.triton.gdn_qkv_split import ( + fused_qkv_split_gdn_prefill, +) +from tokenspeed_kernel.ops.attention.triton.linear.chunk_delta_h import ( + CHUNK_SIZE as FLA_CHUNK_SIZE, +) +from tokenspeed_kernel.ops.attention.triton.linear.index import ( + set_total_chunks_hint, + set_total_chunks_hint_uniform, +) + +from tokenspeed.runtime.configs.paged_cache_spec import LINEAR_ATTENTION +from tokenspeed.runtime.execution.breakable_cuda_graph import ( + break_point, + current_forward_ctx, + scrub_padding_tail, +) +from tokenspeed.runtime.execution.forward_batch_info import ForwardMode +from tokenspeed.runtime.layers.attention.backends.base import ( + AttentionBackend, + init_backend_cuda_graph_state, +) +from tokenspeed.runtime.layers.attention.linear.causal_conv1d import ( + causal_conv1d_fn, + causal_conv1d_update, +) +from tokenspeed.runtime.layers.attention.linear.fused_sigmoid_gating_recurrent import ( + fused_sigmoid_gating_delta_rule_update, +) +from tokenspeed.runtime.layers.attention.linear.gdn import fused_gdn_gating +from tokenspeed.runtime.layers.attention.linear.mamba_state_scatter_triton import ( + fused_mamba_state_copy, +) + +if TYPE_CHECKING: + from tokenspeed.runtime.layers.attention.configs.base import BaseAttnConfig + from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool + from tokenspeed.runtime.layers.paged_attention import PagedAttention + +# Flat KV-cache group id carrying GDN/mamba2 state pages. +_STATE_GROUP_ID = LINEAR_ATTENTION + + +def compute_state_page_indices( + rows: torch.Tensor, + page_size: int, + seq_lens_before: torch.Tensor, + seq_lens_after: torch.Tensor, + *, + validate: bool = True, +) -> tuple[torch.Tensor, torch.Tensor]: + """Dual-index state pages: in = page of position n-1 (0/null when no + history), out = page of the step's last position. rows: [bs, max_pages] + int32 page ids (-1 pad, 0 hole). Within a page in == out (in-place + evolution); crossing a boundary reads the old page and writes the new + one; resuming from a prefix hit reads the claimed snapshot page and + writes the fresh working page. + """ + bs = seq_lens_before.shape[0] + rows = rows[:bs] + before = seq_lens_before.to(torch.int64) + after = seq_lens_after.to(torch.int64) + max_slots = rows.shape[1] + + in_slots = torch.div(before - 1, page_size, rounding_mode="floor").clamp_(min=0) + out_slots = torch.div(after - 1, page_size, rounding_mode="floor") + out_slots_safe = out_slots.clamp(min=0, max=max_slots - 1) + + state_in = rows.gather(1, in_slots.unsqueeze(1)).squeeze(1) + state_in = torch.where(before > 0, state_in, torch.zeros_like(state_in)) + state_out = rows.gather(1, out_slots_safe.unsqueeze(1)).squeeze(1) + + if validate: + if bool((after <= 0).any()): + raise ValueError( + "state paging: seq_lens_after must be >= 1 for every request" + ) + if bool((out_slots >= max_slots).any()): + raise ValueError( + "state paging: out page slot exceeds flat table width " + f"{max_slots} (page_size={page_size})" + ) + if bool((state_in[before > 0] <= 0).any()): + raise ValueError( + "state paging: in page is a pad (-1) or hole (0) for a " + "request with history; reading it would silently resume " + f"from the zero state (flat {_STATE_GROUP_ID!r} table)" + ) + if bool((state_out <= 0).any()): + raise ValueError( + "state paging: out page is a pad (-1) or hole (0); the " + "request's working state page must be present in the flat " + f"{_STATE_GROUP_ID!r} table" + ) + # The <= 0 raise above guarantees every state_out entry is positive. + if torch.unique(state_out).numel() != state_out.numel(): + raise ValueError("state out pages must be unique per batch") + return state_in.to(torch.int32), state_out.to(torch.int32) + + +@dataclass +class MambaForwardMetadata: + query_start_loc: torch.Tensor | None + mamba_cache_indices: torch.Tensor + mamba_output_indices: torch.Tensor | None = None + mamba_req_pool_indices: torch.Tensor | None = None + extend_prefix_lens: torch.Tensor | None = None + extend_seq_lens_cpu: torch.Tensor | None = None + # Pre-computed src/dst indices for extracting Mamba prefix-cache snapshots. + track_ssm_h_src: torch.Tensor | None = None + track_ssm_h_src_fla: torch.Tensor | None = None + track_ssm_h_dst: torch.Tensor | None = None + track_conv_indices: torch.Tensor | None = None + track_ssm_final_src: torch.Tensor | None = None + track_ssm_final_dst: torch.Tensor | None = None + # Flat path (state slabs): dual in/out page indices per request. + state_in_pages: torch.Tensor | None = None + state_out_pages: torch.Tensor | None = None + + +class LayerMappedKVPool: + """Wraps a KV pool to map global layer IDs to internal pool indices. + + For hybrid models, only full attention layers have KV cache. This wrapper + translates global layer IDs (e.g., 3, 7, 11) to pool indices (0, 1, 2). + """ + + def __init__( + self, inner_pool: BaseTokenToKVPool, full_attention_layer_ids: list[int] + ): + self.inner = inner_pool + self.layer_ids = list(full_attention_layer_ids) + self.layer_map = { + global_id: pool_idx + for pool_idx, global_id in enumerate(full_attention_layer_ids) + } + # Expose page_size from inner pool for the scheduler + self.page_size = getattr(inner_pool, "page_size", 1) + + def _map(self, layer_id: int) -> int: + return self.layer_map.get(layer_id, layer_id) + + def set_kv_buffer( + self, + layer: PagedAttention, + out_cache_loc: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor | None, + k_scale: torch.Tensor | None = None, + v_scale: torch.Tensor | None = None, + ): + orig = layer.layer_id + layer.layer_id = self._map(orig) + self.inner.set_kv_buffer(layer, out_cache_loc, k, v, k_scale, v_scale) + layer.layer_id = orig + + def get_kv_buffer(self, layer_id: int): + return self.inner.get_kv_buffer(self._map(layer_id)) + + def get_key_buffer(self, layer_id: int): + return self.inner.get_key_buffer(self._map(layer_id)) + + def get_value_buffer(self, layer_id: int): + return self.inner.get_value_buffer(self._map(layer_id)) + + def __getattr__(self, name): + return getattr(self.inner, name) + + +class SimpleMambaPool: + """Mamba state pool indexed by scheduler-assigned cache slots.""" + + def __init__( + self, + size: int, + num_mamba_layers: int, + conv_state_shape: tuple, + temporal_state_shape: tuple, + conv_dtype: torch.dtype, + ssm_dtype: torch.dtype, + mamba_layer_ids: list[int], + device: str, + page_size: int = 1, + speculative_num_draft_tokens: int = 0, + max_req_pool_size: int = 0, + ): + self.size = size + self.device = device + self.mamba_layer_ids = list(mamba_layer_ids) + self.page_size = page_size + self.mamba_map = {layer_id: i for i, layer_id in enumerate(mamba_layer_ids)} + self.is_kda_cache = False + self.max_req_pool_size = max_req_pool_size + + # Base slots (working + checkpoint) are allocated by C++ scheduler. + # Python-only draft rows live after the scheduler-owned range and are + # addressed by normal row indices in the same tensors. + self.base_size = size + self.speculative_num_draft_tokens = speculative_num_draft_tokens + self.current_input_size = ( + max_req_pool_size + 1 if max_req_pool_size > 0 else size + ) + self.draft_slots_per_req = max(0, speculative_num_draft_tokens - 1) + self.draft_base = size + self.draft_total_slots = self.current_input_size * self.draft_slots_per_req + total_size = size + self.draft_total_slots + self.total_size = total_size + + # Allocate conv state: (num_mamba_layers, total_size, conv_dim, state_len) + self.conv_state = torch.zeros( + num_mamba_layers, + total_size, + *conv_state_shape, + dtype=conv_dtype, + device=device, + ) + # Allocate temporal/SSM state: (num_mamba_layers, total_size, heads, key_dim, val_dim) + self.ssm_state = torch.zeros( + num_mamba_layers, + total_size, + *temporal_state_shape, + dtype=ssm_dtype, + device=device, + ) + + self.mamba_cache = (self.conv_state, self.ssm_state) + self.layer_transfer_counter = None + + self.current_input_indices = torch.full( + (self.current_input_size,), -1, dtype=torch.int32, device=device + ) + + def get_mamba_indices(self, mamba_pool_indices: torch.Tensor) -> torch.Tensor: + """Return mamba cache indices directly (allocated by C++ scheduler).""" + return mamba_pool_indices.to(torch.int32) + + @staticmethod + @torch.compile(dynamic=True) + def _build_mtp_output_indices_kernel( + output_indices: torch.Tensor, + req_pool_indices: torch.Tensor, + working_indices: torch.Tensor, + draft_base: int, + draft_slots_per_req: int, + draft_token_num: int, + ) -> None: + """Fused fill of MTP target-verify output index table. + + Inductor fuses the working-column write and the draft-grid write into + as few elementwise kernels as possible. The host-side early returns + (draft_token_num<=0, ``out is None``) are kept in the wrapper. + """ + bs = working_indices.shape[0] + working = working_indices.to(torch.int32) + valid = working >= 0 + output_indices[:, 0] = torch.where(valid, working, -1) + + if draft_token_num > 1 and draft_slots_per_req > 0: + req = req_pool_indices[:bs].to(torch.int32) + steps = torch.arange( + draft_token_num - 1, dtype=torch.int32, device=working.device + ) + draft = draft_base + req[:, None] * draft_slots_per_req + steps[None, :] + output_indices[:, 1:] = torch.where( + valid[:, None] & (req >= 0)[:, None], + draft, + -1, + ) + + def get_mtp_output_indices( + self, + req_pool_indices: torch.Tensor, + working_indices: torch.Tensor, + draft_token_num: int, + out: torch.Tensor | None = None, + ) -> torch.Tensor: + """Build per-request target-verify outputs: [working, draft0, ...].""" + bs = working_indices.shape[0] + if out is not None: + output_indices = out + output_indices.fill_(-1) + else: + output_indices = torch.full( + (bs, draft_token_num), + -1, + dtype=torch.int32, + device=working_indices.device, + ) + if draft_token_num <= 0: + return output_indices + + self._build_mtp_output_indices_kernel( + output_indices, + req_pool_indices, + working_indices, + self.draft_base, + self.draft_slots_per_req, + draft_token_num, + ) + return output_indices + + @staticmethod + @torch.compile(dynamic=True) + def _get_current_input_indices_kernel( + req_pool_indices: torch.Tensor, + working_indices: torch.Tensor, + current_input_indices_buf: torch.Tensor, + current_input_size: int, + ) -> torch.Tensor: + """Fused gather + masked-where for the no-COW path.""" + n = working_indices.shape[0] + req = req_pool_indices[:n].to(torch.int32) + working = working_indices.to(torch.int32) + valid = (working >= 0) & (req >= 0) + safe = req.clamp(0, current_input_size - 1).to(torch.int64) + stored = current_input_indices_buf[safe] + current = torch.where(valid & (stored >= 0), stored, working) + current = torch.where(valid, current, torch.full_like(current, -1)) + return current + + @staticmethod + @torch.compile(dynamic=True) + def _get_current_input_indices_with_cow_kernel( + req_pool_indices: torch.Tensor, + working_indices: torch.Tensor, + cow_src_indices: torch.Tensor, + current_input_indices_buf: torch.Tensor, + current_input_size: int, + ) -> torch.Tensor: + """Fused gather + masked-where for the COW path.""" + n = working_indices.shape[0] + req = req_pool_indices[:n].to(torch.int32) + working = working_indices.to(torch.int32) + cow = cow_src_indices[:n].to(torch.int32) + valid = (working >= 0) & (req >= 0) + safe = req.clamp(0, current_input_size - 1).to(torch.int64) + stored = current_input_indices_buf[safe] + current = torch.where(valid & (stored >= 0), stored, working) + current = torch.where(valid, current, torch.full_like(current, -1)) + current = torch.where( + (cow >= 0) & valid & (current == working), + cow, + current, + ) + return current + + def get_current_input_indices( + self, + req_pool_indices: torch.Tensor, + working_indices: torch.Tensor, + cow_src_indices: torch.Tensor | None = None, + ) -> torch.Tensor: + """Return the row each request should read at the start of target verify.""" + if cow_src_indices is None: + return self._get_current_input_indices_kernel( + req_pool_indices, + working_indices, + self.current_input_indices, + self.current_input_size, + ) + return self._get_current_input_indices_with_cow_kernel( + req_pool_indices, + working_indices, + cow_src_indices, + self.current_input_indices, + self.current_input_size, + ) + + def reset_current_inputs( + self, req_pool_indices: torch.Tensor, working_indices: torch.Tensor + ) -> None: + """Mark freshly allocated/reused scheduler slots as canonical.""" + req_pool_indices = req_pool_indices[: working_indices.shape[0]].to(torch.int32) + working_indices = working_indices.to(torch.int32) + self.current_input_indices[req_pool_indices.long()] = working_indices + + @staticmethod + @torch.compile(dynamic=True) + def _update_current_inputs_after_verify_kernel( + req_pool_indices: torch.Tensor, + output_indices: torch.Tensor, + accepted_lengths: torch.Tensor, + current_input_indices: torch.Tensor, + max_col: int, + ) -> None: + """Fused gather-scatter for the after-verify input pointer update. + + Inductor fuses clamp/arange/sub/dtype-convert into a single elementwise + kernel; the gather and the in-place scatter on ``current_input_indices`` + each remain a single kernel. All tensors stay on GPU; no host sync. + """ + n = accepted_lengths.shape[0] + req = req_pool_indices[:n].to(torch.int64) + idx = (accepted_lengths.clamp(min=1, max=max_col) - 1).to(torch.int64) + rows = torch.arange(n, device=accepted_lengths.device, dtype=torch.int64) + selected = output_indices[rows, idx].to(torch.int32) + current_input_indices[req] = selected + + def update_current_inputs_after_verify( + self, + req_pool_indices: torch.Tensor, + output_indices: torch.Tensor, + accepted_lengths: torch.Tensor, + ) -> None: + if output_indices is None or output_indices.numel() == 0: + return + self._update_current_inputs_after_verify_kernel( + req_pool_indices, + output_indices, + accepted_lengths, + self.current_input_indices, + output_indices.shape[1], + ) + + def register_layer_transfer_counter(self, layer_transfer_counter): + self.layer_transfer_counter = layer_transfer_counter + + def get_mamba_params(self, layer_id: int): + """Return per-layer cache slices.""" + internal_idx = self.mamba_map[layer_id] + if self.layer_transfer_counter is not None: + self.layer_transfer_counter.wait_until(internal_idx) + return [self.mamba_cache[i][internal_idx] for i in range(len(self.mamba_cache))] + + def get_mamba_params_all_layers(self): + """Return all layers for all cache components.""" + if self.layer_transfer_counter is not None: + self.layer_transfer_counter.wait_until(self.conv_state.shape[0] - 1) + return [self.mamba_cache[i] for i in range(len(self.mamba_cache))] + + def get_contiguous_buf_infos(self): + """Return per-layer mamba cache buffers for disaggregated transfer.""" + data_ptrs = [] + data_lens = [] + item_lens = [] + for cache in self.mamba_cache: + for layer_id in range(cache.shape[0]): + layer_cache = cache[layer_id] + data_ptrs.append(layer_cache.data_ptr()) + data_lens.append(layer_cache.nbytes) + item_lens.append(layer_cache[0].nbytes) + return data_ptrs, data_lens, item_lens + + def get_contiguous_buf_unit_lens(self): + unit_lens = [] + for cache in self.mamba_cache: + for layer_id in range(cache.shape[0]): + layer_cache = cache[layer_id] + unit_lens.append(layer_cache[0, 0].nbytes) + return unit_lens + + def get_contiguous_buf_layer_ids(self): + """Return global layer ids aligned with get_contiguous_buf_infos().""" + return self.mamba_layer_ids * len(self.mamba_cache) + + +class MambaAttnBackend(AttentionBackend): + """Attention backend for Mamba/GDN linear attention layers.""" + + def __init__(self, config: BaseAttnConfig): + super().__init__(config) + self.pad_slot_id = -1 + self.forward_metadata: MambaForwardMetadata = None + self.state_indices_list = [] + self.query_start_loc_list = [] + self.cached_cuda_graph_decode_query_start_loc: torch.Tensor = None + self.cached_cuda_graph_verify_query_start_loc: torch.Tensor = None + self.output_indices_list = [] + self.speculative_num_draft_tokens = getattr( + config, "speculative_num_draft_tokens", 0 + ) + self.pool: SimpleMambaPool = None + # Flat path (state slabs on the KV pool, dual in/out page indexing). + self.kv_pool = None + self.flat_state_active = False + self._flat_state_page_size = 1 + self.flat_state_in_list: list[torch.Tensor] = [] + self.flat_state_out_list: list[torch.Tensor] = [] + + def set_pool(self, pool: SimpleMambaPool): + self.pool = pool + + def set_kv_pool(self, kv_pool) -> None: + """Bind the (layer-mapped) KV pool. Flat state paging turns on iff the + pool allocated state slabs AND publishes the state cache group — + publication (paged_cache_spec.publish_paged_cache_groups) is the + upstream signal that flat block tables will actually be delivered + (radix ext / spec decode never publish).""" + self.kv_pool = kv_pool + specs = getattr(kv_pool, "paged_cache_group_specs", ()) + self.flat_state_active = bool(getattr(kv_pool, "state_slabs", None)) and any( + str(spec.group_id) == _STATE_GROUP_ID for spec in specs + ) + self._flat_state_page_size = int(getattr(kv_pool, "page_size", 1)) + + def _flat_state_pages( + self, + bs: int, + seq_lens: torch.Tensor, + forward_mode: ForwardMode, + kwargs: dict, + *, + validate: bool | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """(state_in, state_out) page ids for this forward, from the flat + state-group block table. seq_lens counts the tokens computed AFTER + this forward (decode: q_len 1; extend: prefix + chunk). + + validate: explicit True/False wins; None (the hot-path default) + validates only under TOKENSPEED_FLAT_DEBUG=1 (the checks host-sync). + """ + if validate is None: + validate = os.environ.get("TOKENSPEED_FLAT_DEBUG") == "1" + flat_tables = kwargs.get("flat_block_tables") + if not flat_tables or _STATE_GROUP_ID not in flat_tables: + raise RuntimeError( + "MambaAttnBackend: flat state paging is active (pool " + f"publishes the {_STATE_GROUP_ID!r} group with state slabs) " + "but flat_block_tables is missing the group " + f"(got {sorted(flat_tables) if flat_tables else flat_tables!r})" + ) + rows = flat_tables[_STATE_GROUP_ID] + after = seq_lens[:bs] + if forward_mode.is_decode_or_idle(): + before = after - 1 + else: + extend_prefix_lens = kwargs.get("extend_prefix_lens") + if extend_prefix_lens is not None: + before = extend_prefix_lens[:bs].to( + device=after.device, dtype=after.dtype + ) + else: + before = torch.zeros_like(after) + return compute_state_page_indices( + rows, + self._flat_state_page_size, + before, + after, + validate=validate, + ) + + def reset_current_inputs( + self, req_pool_indices: torch.Tensor, working_indices: torch.Tensor + ): + if self.pool is not None: + self.pool.reset_current_inputs(req_pool_indices, working_indices) + + def init_forward_metadata( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + forward_mode: ForwardMode = ForwardMode.DECODE, + **kwargs, + ): + mamba_pool_indices = kwargs.get("mamba_pool_indices") + if self.pool is None: + # Poolless flat mode: states are addressed via state_in/out_pages; + # every consumer of mamba_cache_indices is pool/radix-gated. + mamba_cache_indices = None + elif mamba_pool_indices is not None: + mamba_cache_indices = self.pool.get_mamba_indices(mamba_pool_indices[:bs]) + else: + mamba_cache_indices = self.pool.get_mamba_indices(req_pool_indices[:bs]) + + is_target_verify = ( + forward_mode.is_decode_or_idle() + and not self.is_draft + and self.spec_num_tokens > 1 + ) + is_draft_extend = ( + forward_mode.is_decode_or_idle() + and self.is_draft + and self.spec_num_tokens > 1 + ) + + mamba_output_indices = None + extend_seq_lens_cpu = None + if is_target_verify: + draft_token_num = int( + kwargs.get("tokens_per_req", self.speculative_num_draft_tokens) + ) + cow_src_indices = kwargs.get("mamba_cow_src_indices") + mamba_input_indices = self.pool.get_current_input_indices( + req_pool_indices[:bs], mamba_cache_indices, cow_src_indices + ) + mamba_output_indices = self.pool.get_mtp_output_indices( + req_pool_indices[:bs], + mamba_cache_indices, + draft_token_num, + ) + mamba_cache_indices = mamba_input_indices + + if forward_mode.is_decode_or_idle() and self.spec_num_tokens == 1: + query_start_loc = torch.arange( + 0, bs + 1, dtype=torch.int32, device=self.device + ) + elif forward_mode.is_extend_or_mixed() or is_target_verify or is_draft_extend: + if is_target_verify or is_draft_extend: + tokens_per_req = kwargs.get( + "tokens_per_req", self.speculative_num_draft_tokens + ) + query_start_loc = torch.arange( + 0, + bs * tokens_per_req + 1, + step=tokens_per_req, + dtype=torch.int32, + device=self.device, + ) + set_total_chunks_hint_uniform(bs, tokens_per_req, query_start_loc) + else: + extend_start_loc = kwargs.get("extend_start_loc") + extend_seq_lens = kwargs.get("extend_seq_lens") + if extend_start_loc is not None and extend_seq_lens is not None: + query_start_loc = torch.empty( + (bs + 1,), dtype=torch.int32, device=self.device + ) + query_start_loc[:bs] = extend_start_loc + query_start_loc[bs] = extend_start_loc[-1] + extend_seq_lens[-1] + extend_seq_lens_cpu = extend_seq_lens[:bs].to( + device="cpu", dtype=torch.int32 + ) + else: + extend_prefix_lens = kwargs.get("extend_prefix_lens") + if extend_prefix_lens is not None: + extend_lens = (seq_lens[:bs] - extend_prefix_lens[:bs]).to( + torch.int32 + ) + else: + # No prefix: all tokens are new + extend_lens = seq_lens[:bs].to(torch.int32) + query_start_loc = torch.zeros( + bs + 1, dtype=torch.int32, device=self.device + ) + torch.cumsum(extend_lens, dim=0, out=query_start_loc[1:]) + extend_seq_lens_cpu = extend_lens.to(device="cpu") + set_total_chunks_hint(extend_seq_lens_cpu, query_start_loc) + else: + raise ValueError(f"Invalid forward mode: {forward_mode=}") + + state_in_pages = None + state_out_pages = None + # Idle/bs==0 forwards carry no requests and never reach the mamba + # forward (router returns early), so no tables are required. + if self.flat_state_active and bs > 0 and not forward_mode.is_idle(): + if is_target_verify or is_draft_extend: + raise RuntimeError( + "flat GDN state paging does not support speculative " + "decoding; the pool must not publish flat groups under " + "spec (TODO(flat+spec))" + ) + state_in_pages, state_out_pages = self._flat_state_pages( + bs, seq_lens, forward_mode, kwargs + ) + + track_ssm_h_src = None + track_ssm_h_src_fla = None + track_ssm_h_dst = None + track_conv_indices = None + track_ssm_final_src = None + track_ssm_final_dst = None + if ( + (forward_mode.is_extend_or_mixed() or is_draft_extend) + and not is_target_verify + # Radix-only prefix snapshot tracking: flat mode snapshots states + # via page claiming (dual-index), and track indices address the + # SimpleMambaPool row space, not the state slabs. + and not self.flat_state_active + ): + extend_prefix_lens_kw = kwargs.get("extend_prefix_lens") + mamba_track_pool_indices = kwargs.get("mamba_track_pool_indices") + if ( + extend_prefix_lens_kw is not None + and mamba_track_pool_indices is not None + ): + prefix = extend_prefix_lens_kw[:bs].to( + dtype=torch.int32, device=self.device + ) + track_indices = mamba_track_pool_indices[:bs].to( + dtype=torch.int32, device=self.device + ) + extend_lens = (seq_lens[:bs] - prefix).to(torch.int32) + checkpoint_mask = (track_indices >= 0) & (mamba_cache_indices >= 0) + + page_size = getattr(self.pool, "page_size", 1) + final_lens = prefix + extend_lens + last_inserted_lens = (final_lens // page_size) * page_size + track_lens = last_inserted_lens - prefix + track_inside = ( + checkpoint_mask & (track_lens > 0) & (track_lens < extend_lens) + ) + track_mask = track_inside & ((track_lens % FLA_CHUNK_SIZE) == 0) + # C++ attaches the checkpoint slot to the last KV page inserted + # for this chunk. When a chunk has an intermediate branch and + # ends exactly on a page boundary, the final state must win. + final_mask = ( + checkpoint_mask + & (final_lens >= page_size) + & ((final_lens % page_size) == 0) + ) + if final_mask.any(): + track_ssm_final_src = mamba_cache_indices[final_mask] + track_ssm_final_dst = track_indices[final_mask] + + if track_mask.any(): + ( + track_ssm_h_src, + track_ssm_h_src_fla, + track_ssm_h_dst, + ) = self._compute_track_ssm_indices( + track_lens, + track_mask, + track_indices, + seq_lens[:bs] - prefix, # extend_seq_lens + ) + track_conv_indices = self._compute_track_conv_indices( + query_start_loc, + track_lens, + track_mask, + ) + + self.forward_metadata = MambaForwardMetadata( + query_start_loc=query_start_loc, + mamba_cache_indices=mamba_cache_indices, + mamba_output_indices=mamba_output_indices, + mamba_req_pool_indices=req_pool_indices[:bs], + extend_prefix_lens=kwargs.get("extend_prefix_lens"), + extend_seq_lens_cpu=extend_seq_lens_cpu, + track_ssm_h_src=track_ssm_h_src, + track_ssm_h_src_fla=track_ssm_h_src_fla, + track_ssm_h_dst=track_ssm_h_dst, + track_conv_indices=track_conv_indices, + track_ssm_final_src=track_ssm_final_src, + track_ssm_final_dst=track_ssm_final_dst, + state_in_pages=state_in_pages, + state_out_pages=state_out_pages, + ) + + def _compute_track_conv_indices( + self, + query_start_loc: torch.Tensor, + track_lens: torch.Tensor, + track_mask: torch.Tensor, + ): + """Compute packed input indices for conv windows at tracked boundaries.""" + conv_state_len = self.pool.conv_state.shape[-1] + lens_m = track_lens[track_mask] + start = query_start_loc[:-1][track_mask] + lens_m - conv_state_len + indices = start.unsqueeze(-1) + torch.arange( + conv_state_len, + device=self.device, + dtype=start.dtype, + ) + return indices.clamp(0, query_start_loc[-1] - 1) + + def _compute_track_ssm_indices( + self, + track_lens: torch.Tensor, + track_mask: torch.Tensor, + mamba_track_indices: torch.Tensor, + extend_seq_lens: torch.Tensor, + ): + """Compute src/dst indices for extracting intermediate SSM states. + + Matching conv windows are gathered separately from packed pre-conv inputs. + """ + # flashinfer ckpts[k] = state after processing chunk k = FLA h[k+1] + num_fi_ckpts = extend_seq_lens // FLA_CHUNK_SIZE + offset = torch.zeros_like(num_fi_ckpts) + offset[1:] = torch.cumsum(num_fi_ckpts[:-1], dim=0) + num_fla_states = (extend_seq_lens - 1) // FLA_CHUNK_SIZE + 1 + fla_offset = torch.zeros_like(num_fla_states) + fla_offset[1:] = torch.cumsum(num_fla_states[:-1], dim=0) + + lens_m = track_lens[track_mask] + offset_m = offset[track_mask] + fla_offset_m = fla_offset[track_mask] + dst_m = mamba_track_indices[track_mask] + + # FLA h[lens//C] = flashinfer ckpts[lens//C - 1]. + # track_mask guarantees lens_m >= FLA_CHUNK_SIZE so lens_m // C >= 1. + track_ssm_h_src = offset_m + (lens_m // FLA_CHUNK_SIZE - 1) + track_ssm_h_src_fla = fla_offset_m + (lens_m // FLA_CHUNK_SIZE) + track_ssm_h_dst = dst_m + + return ( + track_ssm_h_src, + track_ssm_h_src_fla, + track_ssm_h_dst, + ) + + # ---- CUDA graph state ---- + + def init_cuda_graph_state( + self, max_num_tokens: int, seq_lens_buf: torch.Tensor = None + ): + del seq_lens_buf # mamba doesn't use seq_lens_buf. + for i in range(max_num_tokens): + self.state_indices_list.append( + torch.full( + (i + 1,), self.pad_slot_id, dtype=torch.int32, device=self.device + ) + ) + self.query_start_loc_list.append( + torch.empty((i + 2,), dtype=torch.int32, device=self.device) + ) + if self.flat_state_active: + self.flat_state_in_list.append( + torch.full( + (i + 1,), + self.pad_slot_id, + dtype=torch.int32, + device=self.device, + ) + ) + self.flat_state_out_list.append( + torch.full( + (i + 1,), + self.pad_slot_id, + dtype=torch.int32, + device=self.device, + ) + ) + if self.speculative_num_draft_tokens > 0: + self.output_indices_list.append( + torch.full( + (i + 1, self.speculative_num_draft_tokens), + self.pad_slot_id, + dtype=torch.int32, + device=self.device, + ) + ) + self.cached_cuda_graph_decode_query_start_loc = torch.arange( + 0, max_num_tokens + 1, dtype=torch.int32, device=self.device + ) + if self.speculative_num_draft_tokens > 0: + # Need max_num_tokens+1 entries (one per request + sentinel). + # Each entry is request_index * spec_num_draft_tokens. + self.cached_cuda_graph_verify_query_start_loc = torch.arange( + 0, + (max_num_tokens + 1) * self.speculative_num_draft_tokens, + step=self.speculative_num_draft_tokens, + dtype=torch.int32, + device=self.device, + ) + self._qsl_dirty = [False] * max_num_tokens + self._qsl_last_mode = [None] * max_num_tokens + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + forward_mode: ForwardMode, + **kwargs, + ): + is_target_verify = ( + forward_mode.is_decode_or_idle() + and not self.is_draft + and self.spec_num_tokens > 1 + ) + is_draft_extend = ( + forward_mode.is_decode_or_idle() + and self.is_draft + and self.spec_num_tokens > 1 + ) + + if forward_mode.is_decode_or_idle() and self.spec_num_tokens == 1: + self.query_start_loc_list[bs - 1].copy_( + self.cached_cuda_graph_decode_query_start_loc[: bs + 1] + ) + elif is_target_verify or is_draft_extend: + self.query_start_loc_list[bs - 1].copy_( + self.cached_cuda_graph_verify_query_start_loc[: bs + 1] + ) + else: + raise ValueError(f"Invalid forward mode: {forward_mode=}") + + mamba_pool_indices = kwargs.get("mamba_pool_indices") + # Reuse the pre-allocated [bs]-length buffer as mamba_indices so the + # capture path matches the replay path: zero allocation, single write. + padded_mamba_indices = self.state_indices_list[bs - 1] + if self.pool is None: + # Poolless flat mode: the buffer stays all pad_slot_id (consumers + # are pool/radix-gated; states travel via state_in/out_pages). + padded_mamba_indices.fill_(self.pad_slot_id) + elif mamba_pool_indices is not None: + padded_mamba_indices[:bs].copy_( + self.pool.get_mamba_indices(mamba_pool_indices[:bs]) + ) + else: + padded_mamba_indices[:bs].copy_( + self.pool.get_mamba_indices(req_pool_indices[:bs]) + ) + mamba_output_indices = None + if is_target_verify: + cow_src_indices = kwargs.get("mamba_cow_src_indices") + mamba_input_indices = self.pool.get_current_input_indices( + req_pool_indices[:bs], padded_mamba_indices, cow_src_indices + ) + mamba_output_indices = self.output_indices_list[bs - 1] + self.pool.get_mtp_output_indices( + req_pool_indices[:bs], + padded_mamba_indices, + self.speculative_num_draft_tokens, + out=mamba_output_indices, + ) + padded_mamba_indices.copy_(mamba_input_indices) + state_in_pages = None + state_out_pages = None + if self.flat_state_active: + # Real tables only arrive at replay; capture binds the persistent + # buffers (all pad_slot_id: kernels skip reads/writes at capture, + # so state slab rows are never dirtied by the capture pass). + if is_target_verify or is_draft_extend: + raise RuntimeError( + "flat GDN state paging: CUDA-graph capture supports " + "plain decode only (flat+spec unsupported)" + ) + flat_ids = kwargs.get("flat_cache_group_ids", ()) + if _STATE_GROUP_ID not in flat_ids: + raise RuntimeError( + "flat GDN state paging: capture is missing the " + f"{_STATE_GROUP_ID!r} flat cache group id " + f"(got {tuple(flat_ids)!r})" + ) + state_in_pages = self.flat_state_in_list[bs - 1] + state_out_pages = self.flat_state_out_list[bs - 1] + state_in_pages.fill_(self.pad_slot_id) + state_out_pages.fill_(self.pad_slot_id) + self._qsl_dirty[bs - 1] = False + self._qsl_last_mode[bs - 1] = (forward_mode, self.spec_num_tokens > 1) + self.forward_metadata = MambaForwardMetadata( + query_start_loc=self.query_start_loc_list[bs - 1], + mamba_cache_indices=self.state_indices_list[bs - 1], + mamba_output_indices=mamba_output_indices, + mamba_req_pool_indices=req_pool_indices[:bs], + state_in_pages=state_in_pages, + state_out_pages=state_out_pages, + ) + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + forward_mode: ForwardMode = None, + req_to_page: torch.Tensor = None, + **kwargs, + ): + num_padding = kwargs.get("num_padding", 0) + mamba_pool_indices = kwargs.get("mamba_pool_indices") + + real_bs = bs - num_padding + req_pool_indices = req_pool_indices[:bs] + + # Reuse the pre-allocated [bs]-length buffer as the padded mamba_indices + # so downstream ops (get_mtp_output_indices, get_current_input_indices) + # see the full-batch shape with padding rows already set to -1. + # Zero extra allocations on this hot path. + padded_mamba_indices = self.state_indices_list[bs - 1] + if self.pool is None: + # Poolless flat mode: the buffer stays all pad_slot_id (consumers + # are pool/radix-gated; states travel via state_in/out_pages). + padded_mamba_indices.fill_(self.pad_slot_id) + else: + if mamba_pool_indices is not None: + padded_mamba_indices[:real_bs].copy_( + self.pool.get_mamba_indices(mamba_pool_indices[:real_bs]) + ) + else: + padded_mamba_indices[:real_bs].copy_( + self.pool.get_mamba_indices(req_pool_indices[:real_bs]) + ) + if num_padding > 0: + padded_mamba_indices[real_bs:].fill_(-1) + + is_target_verify = ( + forward_mode is not None + and forward_mode.is_decode_or_idle() + and not self.is_draft + and self.spec_num_tokens > 1 + ) + is_draft_extend = ( + forward_mode is not None + and forward_mode.is_decode_or_idle() + and self.is_draft + and self.spec_num_tokens > 1 + ) + + mamba_output_indices = None + if is_target_verify: + cow_src_indices = kwargs.get("mamba_cow_src_indices") + mamba_input_indices = self.pool.get_current_input_indices( + req_pool_indices, padded_mamba_indices, cow_src_indices + ) + mamba_output_indices = self.output_indices_list[bs - 1] + self.pool.get_mtp_output_indices( + req_pool_indices, + padded_mamba_indices, + self.speculative_num_draft_tokens, + out=mamba_output_indices, + ) + # mamba_input_indices already encodes padding via padded_mamba_indices. + padded_mamba_indices.copy_(mamba_input_indices) + + if num_padding == 0: + need_copy = self._qsl_dirty[bs - 1] or self._qsl_last_mode[bs - 1] != ( + forward_mode, + self.spec_num_tokens > 1, + ) + if need_copy: + if forward_mode.is_decode_or_idle() and self.spec_num_tokens == 1: + self.query_start_loc_list[bs - 1].copy_( + self.cached_cuda_graph_decode_query_start_loc[: bs + 1] + ) + elif is_target_verify or is_draft_extend: + self.query_start_loc_list[bs - 1].copy_( + self.cached_cuda_graph_verify_query_start_loc[: bs + 1] + ) + self._qsl_dirty[bs - 1] = False + self._qsl_last_mode[bs - 1] = (forward_mode, self.spec_num_tokens > 1) + else: + if forward_mode.is_decode_or_idle() and self.spec_num_tokens == 1: + self.query_start_loc_list[bs - 1][:real_bs].copy_( + self.cached_cuda_graph_decode_query_start_loc[:real_bs] + ) + self.query_start_loc_list[bs - 1][real_bs:].fill_(real_bs) + elif is_target_verify or is_draft_extend: + self.query_start_loc_list[bs - 1][:real_bs].copy_( + self.cached_cuda_graph_verify_query_start_loc[:real_bs] + ) + self.query_start_loc_list[bs - 1][real_bs:].fill_( + real_bs * self.speculative_num_draft_tokens + ) + else: + raise ValueError(f"Invalid forward mode: {forward_mode=}") + self._qsl_dirty[bs - 1] = True + self._qsl_last_mode[bs - 1] = (forward_mode, self.spec_num_tokens > 1) + + state_in_pages = None + state_out_pages = None + if self.flat_state_active: + # Decode-only (q_len == 1): before = seq_lens - 1. Padded rows + # (zero table rows, seq_lens 1) are overwritten with pad_slot_id + # below, so validation is skipped to avoid a host sync. + state_in, state_out = self._flat_state_pages( + bs, + seq_lens, + forward_mode, + kwargs, + validate=False, + ) + state_in_pages = self.flat_state_in_list[bs - 1] + state_out_pages = self.flat_state_out_list[bs - 1] + state_in_pages[:real_bs].copy_(state_in[:real_bs]) + state_out_pages[:real_bs].copy_(state_out[:real_bs]) + if num_padding > 0: + state_in_pages[real_bs:].fill_(self.pad_slot_id) + state_out_pages[real_bs:].fill_(self.pad_slot_id) + + self.forward_metadata = MambaForwardMetadata( + query_start_loc=self.query_start_loc_list[bs - 1], + mamba_cache_indices=self.state_indices_list[bs - 1], + mamba_output_indices=mamba_output_indices, + mamba_req_pool_indices=req_pool_indices, + state_in_pages=state_in_pages, + state_out_pages=state_out_pages, + ) + + def get_cuda_graph_seq_len_fill_value(self): + return 1 + + # ---- Forward ---- + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: PagedAttention, + out_cache_loc: torch.Tensor, + token_to_kv_pool, + bs: int, + save_kv_cache: bool = True, + **kwargs, + ): + # Multi-token decode (target verify or drafter compound) reuses + # the multi-token kernel path in forward_extend. `q` is None for + # hybrid linear-attn layers; the token count comes from mixed_qkv. + q_len_per_req = kwargs["mixed_qkv"].shape[0] // bs if bs > 0 else 1 + if q_len_per_req > 1: + return self.forward_extend( + q, + k, + v, + layer, + out_cache_loc, + token_to_kv_pool, + bs, + forward_mode=ForwardMode.DECODE, + save_kv_cache=save_kv_cache, + **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"] + + query_start_loc = self.forward_metadata.query_start_loc + cache_indices = self.forward_metadata.mamba_cache_indices + state_in_pages = self.forward_metadata.state_in_pages + state_out_pages = self.forward_metadata.state_out_pages + use_flat = state_in_pages is not None + + if use_flat: + # Dual-index: read the page holding position n-1, write the page + # holding position n (in == out within a page; a boundary crossing + # reads the old page and writes the new one). pad_slot_id rows + # (graph padding) are skipped by both kernels. + conv_states, ssm_states = self.kv_pool.get_state_buffers(layer_id) + read_indices = state_in_pages + else: + conv_states, ssm_states, *rest = self.pool.get_mamba_params(layer_id) + read_indices = cache_indices + + mixed_qkv = causal_conv1d_update( + mixed_qkv, + conv_states, + conv_weights, + bias, + activation, + conv_state_indices=read_indices, + output_state_indices=(state_out_pages.view(-1, 1) if use_flat else None), + ) + + query, key, value = torch.split( + mixed_qkv, + [ + key_dim // attn_tp_size, + key_dim // attn_tp_size, + value_dim // attn_tp_size, + ], + dim=-1, + ) + 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=read_indices, + cu_seqlens=query_start_loc, + use_qk_l2norm_in_kernel=True, + softplus_beta=1.0, + softplus_threshold=20.0, + # Flat: don't write back to the (possibly shared) in page; the + # post-step state lands on the out page instead. + disable_state_update=use_flat, + output_state_indices=state_out_pages if use_flat else None, + ) + return core_attn_out + + def forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: PagedAttention, + out_cache_loc: torch.Tensor, + token_to_kv_pool, + bs: int, + forward_mode: ForwardMode, + 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"] + + # `q` is None for hybrid linear-attn layers; the token count comes + # from seq_len carried in kwargs. + q_len_per_req = seq_len // bs if bs > 0 else 1 + is_target_verify = ( + forward_mode is not None + and forward_mode.is_decode_or_idle() + and not self.is_draft + and q_len_per_req > 1 + ) + + query_start_loc = self.forward_metadata.query_start_loc + cache_indices = self.forward_metadata.mamba_cache_indices + + if is_target_verify: + draft_token_num = kwargs.get( + "draft_token_num", self.speculative_num_draft_tokens + ) + conv_states, ssm_states = self.pool.get_mamba_params(layer_id) + output_indices = self.forward_metadata.mamba_output_indices + + batch_size = seq_len // draft_token_num + # shouldn't use contiguous here, because causal_conv1d_update + # support input non-contiguous + mixed_qkv_reshaped = mixed_qkv.view( + batch_size, draft_token_num, -1 + ).transpose(1, 2) + mixed_qkv_processed = causal_conv1d_update( + mixed_qkv_reshaped, + conv_states, + conv_weights, + bias, + activation, + conv_state_indices=cache_indices[:batch_size], + output_state_indices=output_indices[:batch_size], + ) + # needn't contiguous here. + mixed_qkv = mixed_qkv_processed.transpose(1, 2).view(seq_len, -1) + else: + state_in_pages = self.forward_metadata.state_in_pages + state_out_pages = self.forward_metadata.state_out_pages + use_flat = state_in_pages is not None + if use_flat: + conv_states, ssm_states = self.kv_pool.get_state_buffers(layer_id) + state_in_long = state_in_pages.to(torch.int64) + state_out_long = state_out_pages.to(torch.int64) + # Seed the out page's conv window from the in page (identity + # within a page; page 0 is the all-zero null page), then run + # the conv read+write entirely on the out page so a shared + # snapshot in-page is never written. + conv_states[state_out_long] = conv_states[state_in_long] + conv_cache_indices = state_out_pages + else: + conv_states, ssm_states = self.pool.get_mamba_params(layer_id) + conv_cache_indices = cache_indices + extend_prefix_lens = kwargs.get("extend_prefix_lens") + if extend_prefix_lens is None: + extend_prefix_lens = self.forward_metadata.extend_prefix_lens + extend_seq_lens_cpu = kwargs.get("extend_seq_lens_cpu") + if extend_seq_lens_cpu is None: + extend_seq_lens_cpu = self.forward_metadata.extend_seq_lens_cpu + has_initial_states = ( + extend_prefix_lens > 0 if extend_prefix_lens is not None else None + ) + need_h_track = ( + self.forward_metadata.track_ssm_h_src is not None + and self.forward_metadata.track_ssm_h_src.numel() > 0 + ) + + # Zero padded rows so garbage can't reach recurrent state (see scrub_padding_tail). + if extend_seq_lens_cpu is not None: + ntok = int(sum(int(x) for x in extend_seq_lens_cpu)) + scrub_padding_tail(ntok, mixed_qkv, a, b) + + mixed_qkv_t = mixed_qkv.transpose(0, 1) + if need_h_track: + if self.forward_metadata.track_conv_indices is None: + raise RuntimeError( + "Missing conv indices for intermediate mamba track" + ) + conv_states[self.forward_metadata.track_ssm_h_dst] = mixed_qkv_t[ + :, self.forward_metadata.track_conv_indices + ].transpose(0, 1) + + mixed_qkv = causal_conv1d_fn( + mixed_qkv_t, + conv_weights, + bias, + activation=activation, + conv_states=conv_states, + has_initial_state=has_initial_states, + cache_indices=conv_cache_indices, + query_start_loc=query_start_loc, + seq_lens_cpu=extend_seq_lens_cpu, + ).transpose(0, 1)[:seq_len] + + key_split_dim = key_dim // attn_tp_size + value_split_dim = value_dim // attn_tp_size + num_heads = key_split_dim // head_k_dim + num_value_heads = value_split_dim // head_v_dim + + query, key, value = fused_qkv_split_gdn_prefill( + mixed_qkv, + num_q_heads=num_heads, + num_k_heads=num_heads, + num_v_heads=num_value_heads, + head_q=head_k_dim, + head_k=head_k_dim, + head_v=head_v_dim, + ) + + if is_target_verify: + draft_token_num = kwargs.get( + "draft_token_num", self.speculative_num_draft_tokens + ) + 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, + # target_verify specific parameters + disable_state_update=True, + output_state_indices=self.forward_metadata.mamba_output_indices, + ) + else: + beta = b.sigmoid() + g = fused_gdn_gating(A_log, a, dt_bias) + g = g.unsqueeze(0) + beta = beta.unsqueeze(0) + + recurrent_state = ssm_states[state_in_long if use_flat else cache_indices] + need_final_track = ( + self.forward_metadata.track_ssm_final_src is not None + and self.forward_metadata.track_ssm_final_src.numel() > 0 + ) + + fi_h_checkpoints = None + h_src = None + if need_h_track: + gdn_result = gdn_chunk_prefill( + query, + key, + value, + g, + beta, + scale=head_k_dim**-0.5, + initial_state=recurrent_state, + cu_seqlens=query_start_loc, + qk_l2norm=True, + output_final_state=True, + output_h=True, + ) + core_attn_out = gdn_result.out + last_recurrent_state = gdn_result.final_state + if gdn_result.h is None: + raise RuntimeError( + "gdn_chunk_prefill(output_h=True) must return checkpoints" + ) + if gdn_result.h_layout is GdnCheckpointLayout.FLASHINFER: + fi_h_checkpoints = gdn_result.h + h_src = self.forward_metadata.track_ssm_h_src + elif gdn_result.h_layout is GdnCheckpointLayout.FLA: + fi_h_checkpoints = gdn_result.h.squeeze(0) + h_src = self.forward_metadata.track_ssm_h_src_fla + else: + raise RuntimeError( + "gdn_chunk_prefill(output_h=True) returned unsupported " + f"checkpoint layout {gdn_result.h_layout}" + ) + else: + gdn_result = gdn_chunk_prefill( + query, + key, + value, + g, + beta, + scale=head_k_dim**-0.5, + initial_state=recurrent_state, + cu_seqlens=query_start_loc, + qk_l2norm=True, + output_final_state=True, + output_h=False, + ) + core_attn_out = gdn_result.out + last_recurrent_state = gdn_result.final_state + last_recurrent_state = last_recurrent_state.to(ssm_states.dtype, copy=False) + ssm_states[state_out_long if use_flat else cache_indices] = ( + last_recurrent_state + ) + + if need_h_track: + ssm_states[self.forward_metadata.track_ssm_h_dst] = fi_h_checkpoints[ + h_src + ].to(ssm_states.dtype, copy=False) + + if need_final_track: + fused_mamba_state_copy( + conv_states, + self.forward_metadata.track_ssm_final_src, + self.forward_metadata.track_ssm_final_dst, + single_layer=True, + ) + fused_mamba_state_copy( + ssm_states, + self.forward_metadata.track_ssm_final_src, + self.forward_metadata.track_ssm_final_dst, + single_layer=True, + ) + + return core_attn_out + + +class HybridLinearAttnBackend(AttentionBackend): + """Hybrid backend that routes between full attention and linear attention by layer ID.""" + + # Both sub-backends consume flat per-group tables (MHA: KV pages; mamba: + # dual-index state pages). Safety comes from the publication rule + # (paged_cache_spec.publish_paged_cache_groups): a radix ext or spec + # decode never publishes groups, so no tables (and no flat capture + # buffers) exist on those paths. + uses_flat_cache_groups: bool = True + + def __init__( + self, + full_attn_backend: AttentionBackend, + linear_attn_backend: MambaAttnBackend, + full_attn_layers: list[int], + ): + self.device = full_attn_backend.device + self.full_attn_layers = set(full_attn_layers) + self.full_attn_backend = full_attn_backend + self.linear_attn_backend = linear_attn_backend + + def _backends(self): + return [self.full_attn_backend, self.linear_attn_backend] + + def _backend_for_layer(self, layer_id: int) -> AttentionBackend: + if self.linear_attn_backend is None or layer_id in self.full_attn_layers: + return self.full_attn_backend + return self.linear_attn_backend + + _MAMBA_KWARGS = frozenset( + { + "mamba_pool_indices", + "mamba_cow_src_indices", + "mamba_branching_seqlens", + "mamba_track_pool_indices", + } + ) + + @staticmethod + def _split_mamba_kwargs(kwargs: dict) -> tuple[dict, dict]: + mamba_kw = {} + common_kw = {} + for k, v in kwargs.items(): + if k in HybridLinearAttnBackend._MAMBA_KWARGS: + mamba_kw[k] = v + else: + common_kw[k] = v + return common_kw, mamba_kw + + # ---- Metadata delegation ---- + + def init_forward_metadata(self, *args, **kwargs): + common_kw, mamba_kw = self._split_mamba_kwargs(kwargs) + self.full_attn_backend.init_forward_metadata(*args, **common_kw) + self.linear_attn_backend.init_forward_metadata(*args, **common_kw, **mamba_kw) + + def init_cuda_graph_state(self, max_bs: int, seq_lens_buf: torch.Tensor, **kwargs): + # kwargs (e.g. paged_cache_group_specs, so the full backend sheds + # state-family groups) are forwarded through the shared signature + # filter: the full backend is user-selectable and may have a narrow + # signature (e.g. TRTLLM MHA takes only (max_bs, seq_lens_buf)), and + # the mamba backend keeps its narrow signature today. + init_backend_cuda_graph_state( + self.full_attn_backend, max_bs, seq_lens_buf, **kwargs + ) + init_backend_cuda_graph_state( + self.linear_attn_backend, max_bs, seq_lens_buf, **kwargs + ) + + def register_step_counter(self, step_counter): + # Hybrid layerwise transfer needs one global step per model layer, + # including both full-attention and mamba layers. Record steps in this + # wrapper instead of in child backends to avoid double counting. + self.step_counter = step_counter + + def init_forward_metadata_capture_cuda_graph(self, *args, **kwargs): + common_kw, mamba_kw = self._split_mamba_kwargs(kwargs) + self.full_attn_backend.init_forward_metadata_capture_cuda_graph( + *args, **common_kw + ) + self.linear_attn_backend.init_forward_metadata_capture_cuda_graph( + *args, **common_kw, **mamba_kw + ) + + def init_forward_metadata_replay_cuda_graph(self, *args, **kwargs): + common_kw, mamba_kw = self._split_mamba_kwargs(kwargs) + self.full_attn_backend.init_forward_metadata_replay_cuda_graph( + *args, **common_kw + ) + self.linear_attn_backend.init_forward_metadata_replay_cuda_graph( + *args, **common_kw, **mamba_kw + ) + + def support_kv_cache_prewrite( + self, forward_mode: ForwardMode | None = None + ) -> bool: + return self.full_attn_backend.support_kv_cache_prewrite(forward_mode) + + # ---- Forward dispatch ---- + + @break_point + def forward( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: PagedAttention, + out_cache_loc, + token_to_kv_pool, + forward_mode: ForwardMode, + bs: int, + save_kv_cache: bool = True, + record_kv_cache: bool | None = None, + **kwargs, + ): + """Dispatch one layer to its full-attention or GDN backend (the break point). + + Overrides the base forward, so it carries its own ``@break_point``; + the frozen capture-time scalars (forward_mode/bs) are re-read from the + ambient ctx (semantics: see breakable_cuda_graph). The GDN scan's + batched [1, T, Hv, D] output is collapsed to z-shaped [T, Hv, D]. + """ + if forward_mode is None: + return super().forward( + q, + k, + v, + layer, + out_cache_loc, + token_to_kv_pool, + forward_mode, + bs, + save_kv_cache, + record_kv_cache=record_kv_cache, + **kwargs, + ) + + # Frozen capture-time scalars, re-read live (see docstring); no-op in eager. + amb = current_forward_ctx() + if amb is not None: + forward_mode = amb.forward_mode + bs = amb.bs + + if 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) + + layer_id = layer.layer_id if layer else kwargs["layer_id"] + backend = self._backend_for_layer(layer_id) + + # See AttentionBackend.forward for the record_kv_cache contract; the step + # is recorded in this wrapper (not the child backends) to keep one step + # per model layer across full-attn + mamba. Idle already returned above. + with self.record_pd_cache_step(forward_mode, save_kv_cache, record_kv_cache): + if forward_mode.is_decode(): + ret = backend.forward_decode( + q, + k, + v, + layer, + out_cache_loc, + token_to_kv_pool, + bs, + save_kv_cache=save_kv_cache, + **kwargs, + ) + else: + ret = backend.forward_extend( + q, + k, + v, + layer, + out_cache_loc, + token_to_kv_pool, + bs, + save_kv_cache=save_kv_cache, + forward_mode=forward_mode, + **kwargs, + ) + # Collapse the GDN scan's batched [1, T, Hv, D] to z-shaped (see docstring). + if ret is not None and ret.dim() == 4: + # Strictly [1, T, Hv, D]: a genuine B>1 must fail loud, not corrupt the handoff. + assert ( + ret.shape[0] == 1 + ), f"GDN scan batched rank expected leading 1, got {ret.shape}" + ret = ret.flatten(0, 1) + return ret + + def forward_decode( + self, q, k, v, layer, out_cache_loc, token_to_kv_pool, bs, **kwargs + ): + layer_id = layer.layer_id if layer else kwargs["layer_id"] + return self._backend_for_layer(layer_id).forward_decode( + q, k, v, layer, out_cache_loc, token_to_kv_pool, bs, **kwargs + ) + + def forward_extend( + self, q, k, v, layer, out_cache_loc, token_to_kv_pool, bs, **kwargs + ): + layer_id = layer.layer_id if layer else kwargs["layer_id"] + return self._backend_for_layer(layer_id).forward_extend( + q, k, v, layer, out_cache_loc, token_to_kv_pool, bs, **kwargs + ) + + def reset_current_inputs(self, *args, **kwargs): + if self.linear_attn_backend is None: + return + if hasattr(self.linear_attn_backend, "reset_current_inputs"): + self.linear_attn_backend.reset_current_inputs(*args, **kwargs) + + def update_mamba_state_after_mtp_verify(self, accepted_length, model): + # mamba_cache_indices are input rows during target-verify. The first + # output row is always the scheduler-owned working slot, so use the + # output index table to update the next-round input pointer. + output_indices = self.linear_attn_backend.forward_metadata.mamba_output_indices + if output_indices is None: + return + req_pool_indices = ( + self.linear_attn_backend.forward_metadata.mamba_req_pool_indices + ) + if req_pool_indices is None: + return + request_number = accepted_length.shape[0] + self.linear_attn_backend.pool.update_current_inputs_after_verify( + req_pool_indices[:request_number], + output_indices[:request_number], + accepted_length, + ) diff --git a/python/tokenspeed/runtime/layers/attention/backends/mha.py b/python/tokenspeed/runtime/layers/attention/backends/mha.py new file mode 100644 index 0000000..5ebcfe4 --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/backends/mha.py @@ -0,0 +1,818 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from functools import partial +from typing import TYPE_CHECKING + +import torch +from tokenspeed_kernel import ( + mha_decode_with_kvcache, + mha_extend_with_kvcache, + mha_plan, + mha_prefill, +) +from tokenspeed_kernel.ops.kvcache.triton import ( + fused_fp8_set_kv_buffer, + gather_page_table_with_padding, +) + +from tokenspeed.runtime.configs.model_config import AttentionArch +from tokenspeed.runtime.execution.breakable_cuda_graph import scrub_padding_tail +from tokenspeed.runtime.execution.forward_batch_info import ForwardMode +from tokenspeed.runtime.layers.attention.backends.base import AttentionBackend +from tokenspeed.runtime.layers.attention.backends.flat_groups import ( + FlatCacheGroupsMixin, +) +from tokenspeed.runtime.layers.attention.configs.mha import MHAConfig +from tokenspeed.runtime.layers.attention.registry import register_backend +from tokenspeed.runtime.layers.attention.utils import build_page_table +from tokenspeed.runtime.utils.common import ceil_div + +if TYPE_CHECKING: + from tokenspeed.runtime.layers.paged_attention import PagedAttention + + +_KERNEL_SOLUTION_BY_BACKEND = { + "mha": None, + "fa3": "fa3", + "fa4": "fa4", + "triton": "triton", + "flashinfer": "flashinfer", +} + + +def _scrub_extend_padding(metadata, q, k, v) -> None: + """Zero the q/k/v rows beyond the real (unpadded) token count under a prefill graph. + + Reads the count from the pinned CPU cu-seqlens mirror (sync-free) and delegates the + zeroing to the shared prefill-graph padding helper. No-op on normal unpadded forwards. + """ + scrub_padding_tail(metadata.cu_extend_seq_lens_cpu[-1], q, k, v) + + +@dataclass(kw_only=True) +class MHAExtendMetadata: + # Device-side metadata: + # - seq_lens: total length after this step + # - extend_seq_lens: length of new tokens + # cu_extend_seq_lens: the cumsum version of extend_seq_lens + # cu_seqlens_kv: the cumsum version of seq_lens + # - extend_prefix_lens: length of the cached prefix tokens + # seq_lens[i] = extend_prefix_lens[i] + extend_seq_lens[i] + # page_table is None on the flat path (per-group page_tables route reads). + page_table: torch.Tensor | None + seq_lens: torch.Tensor + extend_seq_lens: torch.Tensor + cu_extend_seq_lens: torch.Tensor + cu_seqlens_kv: torch.Tensor + extend_prefix_lens: torch.Tensor + extend_seq_lens_cpu: list[int] + cu_extend_seq_lens_cpu: list[int] + max_extend_seq_len: int + max_extend_prefix_len: int = 0 + # Flat per-group page tables (group_id -> [num_reqs, max_pages]); None on + # the single-table path. TODO(radix-removal): drop the single page_table. + page_tables: dict[str, torch.Tensor] | None = None + # Flat per-group KV write locations (group_id -> [num_tokens] int32), + # built with page_tables — same groups, same lifecycle. + out_cache_locs: dict[str, torch.Tensor] | None = None + + +@dataclass(kw_only=True) +class MHADecodeMetadata: + # page_table is None on the flat path (per-group page_tables route reads). + page_table: torch.Tensor | None + seq_lens: torch.Tensor + # Flat per-group tables/write-locs; see MHAExtendMetadata. + page_tables: dict[str, torch.Tensor] | None = None + out_cache_locs: dict[str, torch.Tensor] | None = None + + +class MHAAttnBackend(FlatCacheGroupsMixin, AttentionBackend): + """Standard MHA backend that routes through tokenspeed_kernel attention APIs.""" + + # Unconditional: safety comes from the publication rule + # (paged_cache_spec.publish_paged_cache_groups) plus the replay + # stale-table guard. TODO(radix-removal): drop the flag. + uses_flat_cache_groups: bool = True + + def support_kv_cache_prewrite( + self, forward_mode: ForwardMode | None = None + ) -> bool: + return forward_mode is not None and forward_mode.is_decode() + + def __init__(self, config: MHAConfig): + super().__init__(config) + # Map the selected backend to the corresponding kernel solution string. + backend_name = config.backend_name or "mha" + self.kernel_solution = _KERNEL_SOLUTION_BY_BACKEND[backend_name] + + # Static information needed for metadata construction and kernel dispatch + self.max_context_len = config.context_len + self.page_size = config.page_size + self.max_num_pages = ceil_div(self.max_context_len, self.page_size) + num_q_heads = config.num_attention_heads + num_kv_heads = config.num_kv_heads + self.tp_q_head_num = max(num_q_heads // config.attn_tp_size, 1) + self.tp_kv_head_num = max(num_kv_heads // config.attn_tp_size, 1) + self.head_dim = config.head_dim + self.qkv_dtype = config.dtype + self.kv_cache_dtype = config.kv_cache_dtype + self.is_fp8 = self.kv_cache_dtype in ( + torch.float8_e4m3fn, + torch.float8_e5m2, + ) + self.plan = partial( + mha_plan, + dtype=self.kv_cache_dtype if self.is_fp8 else self.qkv_dtype, + head_dim=self.head_dim, + return_lse=False, + solution=self.kernel_solution, + ) + # DFLASH draft: expand decode metadata to spec_num_tokens rows/request + # (whole block in one decode forward), with uniform non-causal seq_lens. + self.draft_block_decode = bool(getattr(config, "draft_block_decode", False)) + + # Forward metadata is initialized in the runner per forward call + self.forward_decode_metadata: MHADecodeMetadata | None = None + self.forward_extend_metadata: MHAExtendMetadata | None = None + + # ------------------------------------------------------------------ + # Metadata initialization + # ------------------------------------------------------------------ + + def init_forward_metadata( + self, + bs: int, + num_extends: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + req_to_page: torch.Tensor, + forward_mode: ForwardMode, + # Only consumed on the extend/mixed path; decode callers (e.g. the + # DFLASH draft and the cuda-graph wrapper's draft decode init) omit + # them, so they must be optional. + extend_seq_lens: torch.Tensor | None = None, + extend_seq_lens_cpu: torch.Tensor | None = None, + extend_prefix_lens: torch.Tensor | None = None, + extend_prefix_lens_cpu: torch.Tensor | None = None, + flat_block_tables: dict[str, torch.Tensor] | None = None, + **kwargs, + ): + assert not forward_mode.is_mixed(), "mha backend does not support mixed batch" + + seq_lens = seq_lens[:bs] + + flat_page_tables = self._shed_state_groups(flat_block_tables) + flat_out_cache_locs = None + if flat_page_tables: + # Verify keeps [bs]-row tables; only DFLASH expands rows. TODO(flat+dflash). + assert not ( + self.draft_block_decode and self.spec_num_tokens > 1 + ), "flat cache groups are unsupported with DFLASH block decode" + # The flat path routes every read/write through the per-group + # tables; the radix single table would be dead work. + page_table = None + if forward_mode.is_extend_or_mixed(): + assert extend_prefix_lens_cpu is not None + assert extend_seq_lens_cpu is not None + flat_out_cache_locs = self._compute_flat_extend_out_cache_locs( + flat_page_tables, + extend_prefix_lens_cpu[:bs], + extend_seq_lens_cpu[:bs], + self.page_size, + ) + else: + verify_tokens = ( + self.spec_num_tokens + if self.spec_num_tokens > 1 and not self.is_draft + else 1 + ) + flat_out_cache_locs = self._compute_flat_decode_out_cache_locs( + flat_page_tables, + seq_lens, + self.page_size, + verify_tokens, + ) + self._maybe_check_flat_write_locs( + flat_page_tables, flat_out_cache_locs, self.page_size + ) + else: + page_table = build_page_table( + req_pool_indices[:bs], + req_to_page, + self.page_size, + self.max_context_len, + ) + + if forward_mode.is_extend_or_mixed(): + assert extend_seq_lens is not None + assert extend_seq_lens_cpu is not None + assert extend_prefix_lens is not None + assert extend_prefix_lens_cpu is not None + + # Create cumulative sum of the sequence lengths for Q and KV. + extend_seq_lens = extend_seq_lens[:bs] + extend_seq_lens_cpu = [int(x) for x in extend_seq_lens_cpu[:bs].tolist()] + cu_extend_seq_lens = torch.nn.functional.pad( + torch.cumsum(extend_seq_lens, dim=0, dtype=torch.int32), + (1, 0), + ) + cu_extend_seq_lens_cpu = [0] + for length in extend_seq_lens_cpu: + cu_extend_seq_lens_cpu.append(cu_extend_seq_lens_cpu[-1] + length) + cu_seqlens_kv = torch.nn.functional.pad( + torch.cumsum(seq_lens, dim=0, dtype=torch.int32), + (1, 0), + ) + extend_prefix_lens = extend_prefix_lens[:bs] + max_extend_seq_len = max(extend_seq_lens_cpu) + max_extend_prefix_len = int(extend_prefix_lens_cpu[:bs].max().item()) + + self.forward_extend_metadata = MHAExtendMetadata( + page_table=page_table, + seq_lens=seq_lens, + extend_seq_lens=extend_seq_lens, + cu_extend_seq_lens=cu_extend_seq_lens, + cu_seqlens_kv=cu_seqlens_kv, + extend_prefix_lens=extend_prefix_lens, + extend_seq_lens_cpu=extend_seq_lens_cpu, + cu_extend_seq_lens_cpu=cu_extend_seq_lens_cpu, + max_extend_seq_len=max_extend_seq_len, + max_extend_prefix_len=max_extend_prefix_len, + page_tables=flat_page_tables, + out_cache_locs=flat_out_cache_locs, + ) + + # Drafter step 1+ decodes under an EXTEND/MIXED target; seq_lens + # aliases the drafter's live buffer (pre-written by the wrapper). + if self.is_draft: + self.forward_decode_metadata = MHADecodeMetadata( + page_table=page_table, + seq_lens=seq_lens, + page_tables=flat_page_tables, + out_cache_locs=flat_out_cache_locs, + ) + else: + if self.draft_block_decode and self.spec_num_tokens > 1: + # DFLASH drafts a whole block in one decode forward; the decode + # kernel keys masking off max_seqlen_q, so expand each request + # into spec_num_tokens rows with the SAME full seq_len. That + # makes max_seqlen_q == 1 per row, so every block query attends + # over the entire block (non-causal block-diffusion drafting). + # Target verify keeps the unexpanded multi-query decode path. + expanded_page_table, expanded_seq_lens = ( + self._make_spec_metadata_buffers( + bs, + page_table.device, + ) + ) + self._fill_spec_metadata_uniform( + expanded_page_table, + expanded_seq_lens, + page_table, + seq_lens, + ) + self.forward_decode_metadata = MHADecodeMetadata( + page_table=expanded_page_table, + seq_lens=expanded_seq_lens, + page_tables=flat_page_tables, + out_cache_locs=flat_out_cache_locs, + ) + else: + self.forward_decode_metadata = MHADecodeMetadata( + page_table=page_table, + seq_lens=seq_lens, + page_tables=flat_page_tables, + out_cache_locs=flat_out_cache_locs, + ) + + def init_cuda_graph_state( + self, + max_bs: int, + seq_lens_buf: torch.Tensor, + paged_cache_group_specs: Sequence = (), + **kwargs, + ): + # State-family groups (GDN/mamba pages) belong to the mamba backend; + # learn their ids from the pool's specs so every flat table/loc path + # here (eager, capture, replay) sheds them. + self._learn_flat_state_groups(paged_cache_group_specs) + assert ( + seq_lens_buf.dtype == torch.int32 + and seq_lens_buf.dim() == 1 + and seq_lens_buf.shape[0] >= max_bs + ), ( + f"seq_lens_buf must be int32 with shape[0] >= {max_bs}, " + f"got {seq_lens_buf.dtype} {tuple(seq_lens_buf.shape)}" + ) + + self.cuda_graph_decode_metadata = {} + # Flat per-group persistent buffers, lazily allocated at first + # capture. TODO(radix-removal): parallels cuda_graph_page_table. + # Initialized before the DFLASH early return: replay reads the dict + # unconditionally for the stale-table guard. + self._init_flat_graph_buffers(max_bs) + if self.draft_block_decode and self.spec_num_tokens > 1: + # DFLASH draft block: expand to spec_num_tokens decode rows per + # request (one row per block position), so max_seqlen_q == 1 per row + # and every block query attends over the whole block (non-causal). + self.cuda_graph_page_table, self.cuda_graph_seq_lens = ( + self._make_spec_metadata_buffers(max_bs, self.device) + ) + self.cuda_graph_page_table.zero_() + # seq_lens are filled from the live draft length inside the captured + # graph; seed a valid baseline so any pre-broadcast read stays in range. + self.cuda_graph_seq_lens.fill_(self.spec_num_tokens) + return + self.cuda_graph_page_table = torch.zeros( + (max_bs, self.max_num_pages), dtype=torch.int32, device=self.device + ) + if self.spec_num_tokens > 1 and not self.is_draft: + self.cuda_graph_seq_lens = torch.empty( + (max_bs,), dtype=torch.int32, device=self.device + ) + else: + # Alias controller's seq_lens_buf — backend never mutates it. + self.cuda_graph_seq_lens = seq_lens_buf + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + forward_mode: ForwardMode, + flat_cache_group_ids: tuple[str, ...] = (), + **kwargs, + ): + assert not forward_mode.is_extend_or_mixed() + + # Real tables only arrive at replay: capture lazily allocates + # persistent per-group buffers and records metadata views into them, + # so replay can copy_ fresh data to the graph-recorded addresses. + if flat_cache_group_ids: + # Verify keeps [bs]-row tables + [bs*N] loc views. TODO(flat+dflash). + assert not ( + self.draft_block_decode and self.spec_num_tokens > 1 + ), "flat_cache_group_ids is unsupported with DFLASH block decode" + page_tables, out_cache_locs = self._flat_capture_group_views( + bs, + flat_cache_group_ids, + tokens_per_req=( + self.spec_num_tokens + if self.spec_num_tokens > 1 and not self.is_draft + else 1 + ), + ) + + if self.draft_block_decode and self.spec_num_tokens > 1: + # DFLASH draft block: spec_num_tokens decode rows per request. + expanded_bs = bs * self.spec_num_tokens + metadata = MHADecodeMetadata( + page_table=self.cuda_graph_page_table[:expanded_bs, :], + seq_lens=self.cuda_graph_seq_lens[:expanded_bs], + page_tables=page_tables, + out_cache_locs=out_cache_locs, + ) + # Uniform non-causal seq_lens are written by the drafter inside the + # captured graph (see fill_block_decode_seq_lens); seed a safe + # baseline for the capture run before that op records. + metadata.seq_lens.fill_(self.spec_num_tokens) + else: + metadata = MHADecodeMetadata( + # Flat captures route reads through the per-group tables and + # replay never fills the radix single table, so mirror the + # eager flat path: page_table=None instead of a slice of the + # never-filled zero buffer. + page_table=( + None + if page_tables is not None + else self.cuda_graph_page_table[:bs, :] + ), + seq_lens=self.cuda_graph_seq_lens[:bs], + page_tables=page_tables, + out_cache_locs=out_cache_locs, + ) + if self.spec_num_tokens > 1 and not self.is_draft: + metadata.seq_lens.copy_(seq_lens[:bs].clamp_min(self.spec_num_tokens)) + self.cuda_graph_decode_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, + req_to_page: torch.Tensor, + forward_mode: ForwardMode, + flat_block_tables: dict[str, torch.Tensor] | None = None, + **kwargs, + ): + assert not forward_mode.is_extend_or_mixed() + + # Fail loudly instead of replaying over stale/zero page tables. + self._flat_replay_stale_guard(bs, flat_block_tables) + + # Flat captures read only the per-group buffers; the radix single + # table (cuda_graph_page_table) would be dead work there. + if not self.cuda_graph_flat_page_tables: + gather_page_table_with_padding( + req_to_page=req_to_page, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + out=self.cuda_graph_page_table, + bs=bs, + max_num_pages=self.max_num_pages, + page_size=self.page_size, + dummy_slot=0, + ) + if self.spec_num_tokens > 1 and not self.is_draft: + self.cuda_graph_seq_lens[:bs].copy_(seq_lens[:bs]) + elif self.draft_block_decode: + # DFLASH draft: replicate each request's page table to its + # spec_num_tokens block rows. The block-end seq_lens are filled by + # the drafter inside the captured graph, so they are not touched + # here (they re-derive from the live draft length on every replay). + base_page_table = req_to_page[req_pool_indices[:bs], : self.max_num_pages] + self.cuda_graph_page_table[: bs * self.spec_num_tokens, :].view( + bs, self.spec_num_tokens, self.max_num_pages + ).copy_(base_page_table[:, None, :]) + + # cuda_graph_seq_lens is filled by input prep / the spec copy above. + if flat_block_tables: + self._flat_replay_fill( + bs, + flat_block_tables, + self.cuda_graph_seq_lens, + tokens_per_req=( + self.spec_num_tokens + if self.spec_num_tokens > 1 and not self.is_draft + else 1 + ), + ) + + if bs in self.cuda_graph_decode_metadata: + self.forward_decode_metadata = self.cuda_graph_decode_metadata[bs] + + def fill_block_decode_seq_lens(self, bs: int, block_seq_lens: torch.Tensor) -> None: + """DFLASH: broadcast each request's block-end length to its + spec_num_tokens cuda-graph decode rows (uniform, non-causal). + + Called by the drafter inside the captured graph so that on every replay + the expanded seq_lens re-derive from the live draft length (which is + recomputed in-graph from the target's accept lengths). + + Args: + bs: Number of draft requests. + block_seq_lens: ``[bs]`` per-request block-end lengths + (prefix + spec_num_tokens). + """ + spec = self.spec_num_tokens + self.cuda_graph_seq_lens[: bs * spec].view(bs, spec).copy_( + block_seq_lens[:bs] + .clamp(self.spec_num_tokens, self.max_context_len) + .unsqueeze(1) + ) + + # ------------------------------------------------------------------ + # Forward + # ------------------------------------------------------------------ + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor | None, + v: torch.Tensor | None, + layer: PagedAttention, + out_cache_loc: torch.Tensor, + token_to_kv_pool, + bs: int, + save_kv_cache: bool = False, + **kwargs, + ) -> torch.Tensor: + assert layer.qk_head_dim == layer.v_head_dim + assert (k is None) == (v is None) + has_kv = k is not None + + q = q.view(-1, layer.tp_q_head_num, layer.qk_head_dim) + if has_kv: + k = k.view(-1, layer.tp_k_head_num, layer.qk_head_dim) + v = v.view(-1, layer.tp_v_head_num, layer.v_head_dim) + sinks = kwargs.get("sinks") + + out_cache_loc = self._select_out_cache_loc( + layer, + self.forward_decode_metadata, + out_cache_loc, + prefer_caller=self.is_draft, + ) + + return self._forward_decode( + q, + k, + v, + layer, + out_cache_loc, + token_to_kv_pool, + self.forward_decode_metadata, + save_kv_cache=save_kv_cache, + sinks=sinks, + ) + + def forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: PagedAttention, + out_cache_loc: torch.Tensor, + token_to_kv_pool, + bs: int, + save_kv_cache: bool = False, + **kwargs, + ) -> torch.Tensor: + assert layer.qk_head_dim == layer.v_head_dim + assert (k is None) == (v is None) + assert k is not None + + q = q.view(-1, layer.tp_q_head_num, layer.qk_head_dim) + k = k.view(-1, layer.tp_k_head_num, layer.qk_head_dim) + v = v.view(-1, layer.tp_v_head_num, layer.v_head_dim) + + metadata = self.forward_extend_metadata + sinks = kwargs.get("sinks") + out_cache_loc = self._select_out_cache_loc(layer, metadata, out_cache_loc) + plan = self.plan( + window_left=layer.sliding_window_size, + logit_cap=layer.logit_cap, + sinks=sinks, + ) + + extend_mode = plan.get("extend_mode", "prewrite") + if metadata.max_extend_prefix_len == 0 and extend_mode == "postwrite": + return self._forward_prefill( + q, + k, + v, + layer, + out_cache_loc, + token_to_kv_pool, + metadata, + save_kv_cache, + sinks, + ) + else: + return self._forward_extend( + q, + k, + v, + layer, + out_cache_loc, + token_to_kv_pool, + metadata, + save_kv_cache, + sinks, + ) + + def _forward_prefill( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: PagedAttention, + out_cache_loc: torch.Tensor, + token_to_kv_pool, + metadata: MHAExtendMetadata, + save_kv_cache: bool, + sinks: torch.Tensor | None, + ) -> torch.Tensor: + _scrub_extend_padding(metadata, q, k, v) + # TODO: use a custom kernel to do downcast + if self.is_fp8: + q = q.to(self.kv_cache_dtype) + k = k.to(self.kv_cache_dtype) + v = v.to(self.kv_cache_dtype) + + output = mha_prefill( + q=q, + k=k, + v=v, + cu_seqlens=metadata.cu_extend_seq_lens, + cu_seqlens_cpu=metadata.cu_extend_seq_lens_cpu, + max_seqlen=metadata.max_extend_seq_len, + window_left=layer.sliding_window_size, + logit_cap=layer.logit_cap, + sinks=sinks, + solution=self.kernel_solution, + ) + output = output.reshape(-1, layer.tp_q_head_num * layer.v_head_dim) + if save_kv_cache: + self._save_kv_cache(layer, out_cache_loc, token_to_kv_pool, k, v) + return output + + def _forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor | None, + v: torch.Tensor | None, + layer: PagedAttention, + out_cache_loc: torch.Tensor, + token_to_kv_pool, + metadata: MHAExtendMetadata, + save_kv_cache: bool, + sinks: torch.Tensor | None, + ) -> torch.Tensor: + _scrub_extend_padding(metadata, q, k, v) + if save_kv_cache: + self._save_kv_cache(layer, out_cache_loc, token_to_kv_pool, k, v) + + if self.is_fp8: + q = q.to(self.kv_cache_dtype) + + k_cache, v_cache = self._get_kv_cache(layer, token_to_kv_pool) + output = mha_extend_with_kvcache( + q=q, + cu_seqlens_q=metadata.cu_extend_seq_lens, + cu_seqlens_kv=metadata.cu_seqlens_kv, + k_cache=k_cache, + v_cache=v_cache, + page_table=self._select_page_table(layer, metadata), + cache_seqlens=metadata.seq_lens, + max_seqlen_q=metadata.max_extend_seq_len, + max_seqlen_k=self.max_context_len, + # DFLASH marks its draft attention non-causal so the draft block's + # query positions attend bidirectionally. Every other layer leaves + # the attribute unset, so this stays causal by default. + is_causal=not bool(getattr(layer, "non_causal", False)), + window_left=layer.sliding_window_size, + logit_cap=layer.logit_cap, + sinks=sinks, + solution=self.kernel_solution, + ) + return output.reshape(-1, layer.tp_q_head_num * layer.v_head_dim) + + def _forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor | None, + v: torch.Tensor | None, + layer: PagedAttention, + out_cache_loc: torch.Tensor, + token_to_kv_pool, + metadata: MHADecodeMetadata, + save_kv_cache: bool, + sinks: torch.Tensor | None, + ) -> torch.Tensor: + if save_kv_cache: + self._save_kv_cache(layer, out_cache_loc, token_to_kv_pool, k, v) + + if self.is_fp8: + q = q.to(self.kv_cache_dtype) + + k_cache, v_cache = self._get_kv_cache(layer, token_to_kv_pool) + max_seqlen_q = q.shape[0] // metadata.seq_lens.shape[0] + output = mha_decode_with_kvcache( + q=q, + k_cache=k_cache, + v_cache=v_cache, + page_table=self._select_page_table(layer, metadata), + cache_seqlens=metadata.seq_lens, + window_left=layer.sliding_window_size, + logit_cap=layer.logit_cap, + sinks=sinks, + max_seqlen_k=self.max_context_len, + max_seqlen_q=max_seqlen_q, + solution=self.kernel_solution, + ) + return output.reshape(-1, layer.tp_q_head_num * layer.v_head_dim) + + # ------------------------------------------------------------------ + # Helper methods + # ------------------------------------------------------------------ + + def _save_kv_cache( + self, + layer: PagedAttention, + out_cache_loc: torch.Tensor, + token_to_kv_pool, + k: torch.Tensor | None, + v: torch.Tensor | None, + ) -> None: + if k is None: + return + + if ( + self.kv_cache_dtype == torch.float8_e4m3fn + and k.dtype != torch.float8_e4m3fn + ): + k_cache, v_cache = token_to_kv_pool.get_kv_buffer(layer.layer_id) + fused_fp8_set_kv_buffer( + k=k, + v=v, + k_cache=k_cache, + v_cache=v_cache, + cache_loc=out_cache_loc, + k_scale=layer.k_scale, + v_scale=layer.v_scale, + page_size=self.page_size, + ) + else: + token_to_kv_pool.set_kv_buffer( + layer, + out_cache_loc, + k, + v, + layer.k_scale, + layer.v_scale, + ) + + def _get_kv_cache(self, layer: PagedAttention, token_to_kv_pool): + k_cache = 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 = token_to_kv_pool.get_value_buffer(layer.layer_id).view( + -1, + self.page_size, + layer.tp_v_head_num, + layer.v_head_dim, + ) + return k_cache, v_cache + + def _make_spec_metadata_buffers( + self, + bs: int, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor]: + expanded_bs = bs * self.spec_num_tokens + cuda_graph_page_table = torch.empty( + (expanded_bs, self.max_num_pages), + dtype=torch.int32, + device=device, + ) + cuda_graph_seq_lens = torch.empty( + (expanded_bs,), + dtype=torch.int32, + device=device, + ) + return (cuda_graph_page_table, cuda_graph_seq_lens) + + def _fill_spec_metadata_uniform( + self, + expanded_page_table: torch.Tensor, + expanded_seq_lens: torch.Tensor, + page_table: torch.Tensor, + seq_lens: torch.Tensor, + ): + """Expand spec metadata with a uniform (non-causal) seq_len per row. + + Replicates the full seq_len to all spec_num_tokens rows of a request so + each row decodes with max_seqlen_q == 1 over the whole block. Used by the + DFLASH drafter so every block query attends over the entire block + (non-causal block-diffusion drafting), as opposed to the target's + unexpanded causal multi-query verify path. + """ + bs = seq_lens.shape[0] + spec_num_tokens = self.spec_num_tokens + expanded_page_table = expanded_page_table.view( + bs, spec_num_tokens, self.max_num_pages + ) + expanded_page_table.copy_(page_table[:, None, :]) + # Clamp to max_context_len so the draft decode never asks the attention + # kernel for more than max_num_pages worth of page-table columns. The + # block-end length is prefix + spec_num_tokens, which can exceed + # max_context_len for a request near the context limit; without the + # clamp the kernel reads page_table[:, >= max_num_pages] out of bounds + # (CUDA illegal memory access). Mirrors fill_block_decode_seq_lens on the + # cuda-graph path (this eager path is taken by mixed prefill+decode + # batches even when cuda graphs are enabled). + expanded_seq_lens.view(bs, spec_num_tokens).copy_( + seq_lens.clamp(spec_num_tokens, self.max_context_len)[:, None] + ) + + +for _backend_name in _KERNEL_SOLUTION_BY_BACKEND: + register_backend(_backend_name, {AttentionArch.MHA}, MHAAttnBackend) diff --git a/python/tokenspeed/runtime/layers/attention/backends/mla.py b/python/tokenspeed/runtime/layers/attention/backends/mla.py new file mode 100644 index 0000000..c19bc8f --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/backends/mla.py @@ -0,0 +1,475 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from contextlib import contextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch +from tokenspeed_kernel import mla_decode_with_kvcache, mla_prefill + +from tokenspeed.runtime.configs.model_config import AttentionArch +from tokenspeed.runtime.execution.forward_batch_info import ForwardMode +from tokenspeed.runtime.layers.attention.backends.base import AttentionBackend +from tokenspeed.runtime.layers.attention.chunk import ( + build_chunked_prefill_metadata_arrays, +) +from tokenspeed.runtime.layers.attention.configs.mla import MLAConfig +from tokenspeed.runtime.layers.attention.registry import register_backend +from tokenspeed.runtime.layers.attention.utils import build_page_table +from tokenspeed.runtime.utils.common import ceil_div + +if TYPE_CHECKING: + from tokenspeed.runtime.layers.paged_attention import PagedAttention + + +@dataclass(kw_only=True) +class MLAPrefillMetadata: + # Device-side metadata for explicit Q/K/V MLA prefill and prefix replay. + seq_lens: torch.Tensor + req_pool_indices: torch.Tensor + extend_prefix_lens: torch.Tensor + extend_seq_lens: torch.Tensor + cum_extend_seq_lens: torch.Tensor + # Host-side metadata. + extend_seq_lens_cpu: list[int] + max_extend_seq_len: int + max_extend_prefix_len: int + # Per-prefix-chunk arrays consumed by DeepSeek's chunked prefix replay. + chunked_loop_num: int + chunk_kv_indices_list: list[torch.Tensor] + chunked_seq_len: torch.Tensor + cu_chunked_seq_len: torch.Tensor + max_chunk_len_per_loop: list[int] + + +@dataclass(kw_only=True) +class MLADecodeMetadata: + # num_extends lets mixed batches slice decode requests after extend requests. + num_extends: int + page_table: torch.Tensor + seq_lens: torch.Tensor + + @property + def block_kv_indices(self) -> torch.Tensor: + return self.page_table + + @property + def seq_lens_k(self) -> torch.Tensor: + return self.seq_lens + + +class MLAAttnBackend(AttentionBackend): + """Unified MLA backend routed through tokenspeed_kernel MLA APIs.""" + + def __init__(self, config: MLAConfig): + super().__init__(config) + + self.max_context_len = config.context_len + self.page_size = config.page_size + self.max_num_pages = ceil_div(self.max_context_len, self.page_size) + + 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 = config.kv_cache_dim + self.scaling = config.scaling + self.data_type = config.kv_cache_dtype + self.q_data_type = config.dtype + self.num_local_heads = config.num_attention_heads // config.attn_tp_size + + self.kernel_solution = None + self.forward_decode_metadata: MLADecodeMetadata | None = None + self.forward_prefill_metadata: MLAPrefillMetadata | None = None + self.chunked_prefill_metadata: MLAPrefillMetadata | None = None + self.decode_cuda_graph_metadata: dict[int, MLADecodeMetadata] = {} + self.cuda_graph_page_table: torch.Tensor | None = None + self.cuda_graph_seq_lens: torch.Tensor | None = None + + def init_forward_metadata( + self, + bs: int, + num_extends: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + req_to_page: torch.Tensor, + forward_mode: ForwardMode, + extend_seq_lens: torch.Tensor | None = None, + extend_seq_lens_cpu: torch.Tensor | None = None, + extend_prefix_lens: torch.Tensor | None = None, + extend_prefix_lens_cpu: torch.Tensor | None = None, + **kwargs, + ): + if forward_mode.is_extend_or_mixed(): + self._init_prefill_metadata( + seq_lens=seq_lens[:num_extends], + req_pool_indices=req_pool_indices[:num_extends], + req_to_page=req_to_page, + extend_prefix_lens=extend_prefix_lens[:num_extends], + extend_prefix_lens_cpu=extend_prefix_lens_cpu[:num_extends], + extend_seq_lens=extend_seq_lens[:num_extends], + extend_seq_lens_cpu=extend_seq_lens_cpu[:num_extends], + ) + + if ( + forward_mode.is_decode() + or forward_mode.is_mixed() + or (forward_mode.is_extend() and self.is_draft) + ): + self._init_decode_metadata( + bs=bs, + num_extends=num_extends, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + req_to_page=req_to_page, + ) + + @contextmanager + def override_num_extends(self, num_extends: int): + assert self.forward_decode_metadata is not None + prev = self.forward_decode_metadata.num_extends + self.forward_decode_metadata.num_extends = num_extends + try: + yield + finally: + self.forward_decode_metadata.num_extends = prev + + def _init_prefill_metadata( + self, + seq_lens: torch.Tensor, + req_pool_indices: torch.Tensor, + req_to_page: torch.Tensor, + extend_prefix_lens: torch.Tensor, + extend_prefix_lens_cpu: torch.Tensor, + extend_seq_lens: torch.Tensor, + extend_seq_lens_cpu: torch.Tensor, + ): + extend_seq_lens_cpu_list = [int(x) for x in extend_seq_lens_cpu.tolist()] + cum_extend_seq_lens = torch.zeros( + extend_seq_lens.shape[0] + 1, + device=self.device, + dtype=torch.int32, + ) + torch.cumsum(extend_seq_lens, dim=0, out=cum_extend_seq_lens[1:]) + + max_extend_seq_len = max(extend_seq_lens_cpu_list, default=0) + max_extend_prefix_len = int(extend_prefix_lens_cpu.max().item()) + + ( + chunked_loop_num, + chunk_kv_indices_list, + chunked_seq_len, + cu_chunked_seq_len, + max_chunk_len_per_loop, + ) = build_chunked_prefill_metadata_arrays( + extend_prefix_lens, + extend_prefix_lens_cpu, + req_to_page, + req_pool_indices, + self.page_size, + ) + + metadata = MLAPrefillMetadata( + seq_lens=seq_lens, + req_pool_indices=req_pool_indices, + extend_prefix_lens=extend_prefix_lens, + extend_seq_lens=extend_seq_lens, + cum_extend_seq_lens=cum_extend_seq_lens, + extend_seq_lens_cpu=extend_seq_lens_cpu_list, + max_extend_seq_len=max_extend_seq_len, + max_extend_prefix_len=max_extend_prefix_len, + chunked_loop_num=chunked_loop_num, + chunk_kv_indices_list=chunk_kv_indices_list, + chunked_seq_len=chunked_seq_len, + cu_chunked_seq_len=cu_chunked_seq_len, + max_chunk_len_per_loop=max_chunk_len_per_loop, + ) + self.forward_prefill_metadata = metadata + self.chunked_prefill_metadata = metadata + + def _init_decode_metadata( + self, + bs: int, + num_extends: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + req_to_page: torch.Tensor, + ): + page_table = build_page_table( + req_pool_indices[:bs], + req_to_page, + self.page_size, + self.max_context_len, + ) + self.forward_decode_metadata = MLADecodeMetadata( + num_extends=num_extends, + page_table=page_table, + seq_lens=seq_lens[:bs], + ) + + def init_cuda_graph_state(self, max_bs: int, seq_lens_buf: torch.Tensor): + assert ( + seq_lens_buf.dtype == torch.int32 + and seq_lens_buf.dim() == 1 + and seq_lens_buf.shape[0] >= max_bs + ), ( + f"seq_lens_buf must be int32 with shape[0] >= {max_bs}, " + f"got {seq_lens_buf.dtype} {tuple(seq_lens_buf.shape)}" + ) + self.cuda_graph_page_table = torch.zeros( + (max_bs, self.max_num_pages), dtype=torch.int32, device=self.device + ) + self.cuda_graph_seq_lens = seq_lens_buf + self.decode_cuda_graph_metadata = {} + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + forward_mode: ForwardMode, + ): + if forward_mode.is_extend_or_mixed(): + raise NotImplementedError( + f"mla CUDA graph capture not supported for {forward_mode}" + ) + + metadata = MLADecodeMetadata( + num_extends=0, + page_table=self.cuda_graph_page_table[:bs, :], + seq_lens=self.cuda_graph_seq_lens[:bs], + ) + 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, + forward_mode: ForwardMode = None, + req_to_page: torch.Tensor = None, + **kwargs, + ): + if forward_mode is not None and forward_mode.is_extend_or_mixed(): + raise NotImplementedError( + f"mla CUDA graph replay not supported for {forward_mode}" + ) + + self.cuda_graph_page_table[:bs, : self.max_num_pages].copy_( + req_to_page[req_pool_indices[:bs], : self.max_num_pages] + ) + self.forward_decode_metadata = self.decode_cuda_graph_metadata[bs] + + 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: PagedAttention, + out_cache_loc: torch.Tensor, + token_to_kv_pool, + bs: int, + save_kv_cache: bool = True, + **kwargs, + ) -> torch.Tensor: + # q is absorbed MLA query [T, H, R + D_rope]; k is compressed KV + # [T, 1, R + D_rope]. DeepSeek normally writes cache before this call. + if save_kv_cache: + assert k is not None + token_to_kv_pool.set_mla_kv_buffer( + layer, + out_cache_loc, + k[..., : self.kv_lora_rank], + k[..., self.kv_lora_rank :], + ) + + metadata = self.forward_decode_metadata + assert metadata is not None + num_extends = metadata.num_extends + q_len_per_req = q.shape[0] // bs if bs > 0 else 1 + + if q_len_per_req > 1: + query = q.view(-1, layer.tp_q_head_num, layer.head_dim).unsqueeze(1) + page_table = metadata.page_table[num_extends:].repeat_interleave( + q_len_per_req, dim=0 + ) + cache_seqlens = metadata.seq_lens[num_extends:].repeat_interleave( + q_len_per_req + ) + # Draft catch-up starts from the current draft KV length; target + # verify starts from the final target KV length and backs up. + offset_start = 0 if self.is_draft else 1 - q_len_per_req + offsets = torch.arange( + offset_start, + offset_start + q_len_per_req, + device=cache_seqlens.device, + dtype=cache_seqlens.dtype, + ).repeat(bs) + cache_seqlens = cache_seqlens + offsets + max_seqlen_k = self.max_context_len + else: + query = q.view(bs, -1, layer.tp_q_head_num, layer.head_dim) + page_table = metadata.page_table[num_extends:] + cache_seqlens = metadata.seq_lens[num_extends:] + max_seqlen_k = self.max_context_len + + softmax_scale = layer.scaling + if self.data_type == torch.float8_e4m3fn: + query = query.to(self.data_type) + k_scale = ( + layer.k_scale_float + if getattr(layer, "k_scale_float", None) is not None + else 1.0 + ) + softmax_scale = k_scale * softmax_scale + + kv_cache = token_to_kv_pool.get_key_buffer(layer.layer_id) + if self.data_type != kv_cache.dtype: + kv_cache = kv_cache.to(self.data_type) + kv_cache = kv_cache.view(-1, self.page_size, 1, self.kv_cache_dim) + + result = mla_decode_with_kvcache( + q=query, + kv_cache=kv_cache, + page_table=page_table, + cache_seqlens=cache_seqlens, + max_seqlen_k=max_seqlen_k, + 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, + softmax_scale=softmax_scale, + logit_cap=layer.logit_cap, + solution=self.kernel_solution, + ) + output = self._unwrap_output(result) + return output.reshape(-1, layer.tp_q_head_num * layer.v_head_dim) + + def forward_extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: PagedAttention, + out_cache_loc: torch.Tensor, + token_to_kv_pool, + bs: int, + save_kv_cache: bool = True, + **kwargs, + ) -> torch.Tensor: + if save_kv_cache: + raise NotImplementedError( + "MLA forward_extend cannot derive compressed cache rows from " + "materialized K/V; DeepSeek writes MLA cache in the model path" + ) + + metadata = self.forward_prefill_metadata + assert metadata is not None + if metadata.max_extend_prefix_len > 0: + raise NotImplementedError( + "MLA prefix-cache extend is handled by DeepSeek's chunked " + "prefix replay path via forward_extend_chunked" + ) + + q = q.view(-1, layer.tp_q_head_num, layer.qk_head_dim) + k = k.view(-1, layer.tp_k_head_num, layer.qk_head_dim) + v = v.view(-1, layer.tp_v_head_num, layer.v_head_dim) + result = mla_prefill( + q=q, + k=k, + v=v, + cu_seqlens_q=metadata.cum_extend_seq_lens, + cu_seqlens_kv=metadata.cum_extend_seq_lens, + max_seqlen_q=metadata.max_extend_seq_len, + max_seqlen_kv=metadata.max_extend_seq_len, + softmax_scale=layer.scaling, + seq_lens_kv=metadata.extend_seq_lens, + is_causal=True, + logit_cap=layer.logit_cap, + solution=self.kernel_solution, + ) + output = self._unwrap_output(result) + return output.reshape(-1, layer.tp_q_head_num * layer.v_head_dim) + + def forward_extend_chunked( + self, + q, + k, + v, + scaling, + logits_soft_cap=None, + *, + cum_seq_lens_q, + cum_seq_lens_kv, + max_q_len, + max_kv_len, + seq_lens, + batch_size, + causal, + out: torch.Tensor | None = None, + ): + if causal: + step_counter = getattr(self, "step_counter", None) + if step_counter is not None: + step_counter.record_cache() + + head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + q = q.reshape(-1, self.num_local_heads, head_dim) + k = k.reshape(-1, self.num_local_heads, head_dim) + v = v.reshape(-1, self.num_local_heads, self.v_head_dim) + + if q.dtype == torch.float8_e4m3fn: + k = k.to(torch.float8_e4m3fn) + v = v.to(torch.float8_e4m3fn) + + result = mla_prefill( + q=q, + k=k, + v=v, + cu_seqlens_q=cum_seq_lens_q, + cu_seqlens_kv=cum_seq_lens_kv, + max_seqlen_q=max_q_len, + max_seqlen_kv=max_kv_len, + softmax_scale=scaling, + seq_lens_kv=seq_lens, + is_causal=causal, + logit_cap=logits_soft_cap or 0.0, + return_lse=True, + out=out, + solution=self.kernel_solution, + ) + + if isinstance(result, tuple): + return result[0], result[1] + return result, None + + def _unwrap_output(self, result): + if isinstance(result, tuple): + return result[0] + return result + + +register_backend("mla", {AttentionArch.MLA}, MLAAttnBackend) diff --git a/python/tokenspeed/runtime/layers/attention/backends/tokenspeed_mla.py b/python/tokenspeed/runtime/layers/attention/backends/tokenspeed_mla.py new file mode 100644 index 0000000..7dab12f --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/backends/tokenspeed_mla.py @@ -0,0 +1,563 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +""" +CuteDSL MLA attention backend for TokenSpeed scheduling. + +Uses CuTe DSL JIT-compiled kernels for MLA decode and prefill on Blackwell SM100 GPUs: +- tokenspeed_mla_decode for decode/verify +- tokenspeed_mla_prefill for prefill +""" + +from __future__ import annotations + +import logging +from contextlib import contextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch +import triton +from tokenspeed_kernel.ops.attention.tokenspeed_mla import ( + get_num_sm, + tokenspeed_mla_decode, + tokenspeed_mla_prefill, + warmup_compile_prefill, +) + +from tokenspeed.runtime.configs.model_config import AttentionArch +from tokenspeed.runtime.execution.forward_batch_info import ForwardMode +from tokenspeed.runtime.layers.attention.backends.base import AttentionBackend +from tokenspeed.runtime.layers.attention.backends.trtllm_mla import ( + TRTLLM_BLOCK_CONSTRAINT, + TRTLLMMLAChunkedPrefillMetadata, +) +from tokenspeed.runtime.layers.attention.chunk import ( + build_chunked_prefill_metadata_arrays, +) +from tokenspeed.runtime.layers.attention.configs.mla import MLAConfig +from tokenspeed.runtime.layers.attention.registry import register_backend +from tokenspeed.runtime.utils.env import global_server_args_dict +from tokenspeed.runtime.utils.pdl import pdl_enabled + +if TYPE_CHECKING: + from tokenspeed.runtime.layers.paged_attention import PagedAttention + +logger = logging.getLogger(__name__) + +# CuteDSL decode workspace. The kernel's own `get_workspace_size` formula is +# B * H * q_len * split_kv * (D + 1) * (acc_dtype.width // 8) (bytes) +# and `B * split_kv <= num_SMs`, so a closed-form upper bound (Float32 acc) is +# num_SMs * H * q_len * (D + 1) * 4 +# Buffer is per-device and does NOT need zero-init. +_cutedsl_workspace_buffer: dict[torch.device, torch.Tensor] = {} + +# Initial q_len capacity for the per-device decode workspace. The buffer grows +# on demand before each launch, so larger verify/draft batches are supported. +_CUTEDSL_INITIAL_Q_LEN_CAPACITY = 8 + + +def get_cutedsl_workspace_buffer( + device: torch.device, + num_heads_per_tp: int, + kv_lora_rank: int, + q_len_capacity: int = _CUTEDSL_INITIAL_Q_LEN_CAPACITY, +) -> torch.Tensor: + """Get or grow the per-device CuteDSL workspace buffer.""" + num_sms = get_num_sm(device) + required = num_sms * num_heads_per_tp * q_len_capacity * (kv_lora_rank + 1) * 4 + + existing = _cutedsl_workspace_buffer.get(device) + if existing is None or existing.numel() < required: + _cutedsl_workspace_buffer[device] = torch.empty( + required, dtype=torch.int8, device=device + ) + return _cutedsl_workspace_buffer[device] + + +@dataclass +class CuteDSLMLAPrefillMetadata: + max_seq_len: int + cum_seq_lens: torch.Tensor + seq_lens: torch.Tensor + + +@dataclass +class CuteDSLMLADecodeMetadata: + num_extends: int = 0 + block_kv_indices: torch.Tensor | None = None + max_seq_len_k: int | None = None + seq_lens_k: torch.Tensor | None = None + + +class CuteDSLMLABackend(AttentionBackend): + """CuteDSL MLA attention backend for Blackwell SM100 GPUs. + + Decode uses CuTe DSL JIT-compiled kernels via tokenspeed_mla_decode(). + Prefill uses CuTe DSL FMHA kernel via tokenspeed_mla_prefill(). + """ + + _logged_decode = False + _logged_prefill = False + + def __init__(self, config: MLAConfig): + super().__init__(config) + + self.max_context_len = config.context_len + self.page_size = config.page_size + + # MLA 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 = config.kv_cache_dim + self.scaling = config.scaling + self.data_type = config.kv_cache_dtype + self.q_data_type = config.dtype + + # Workspace buffers — sized from config's num_heads / kv_lora_rank. + num_heads_per_tp = config.num_attention_heads // config.attn_tp_size + self.cutedsl_workspace = get_cutedsl_workspace_buffer( + config.device, num_heads_per_tp, self.kv_lora_rank + ) + + # Pre-compile prefill kernel variants so JIT doesn't run during serving. + # The backend may be constructed once per attention layer (60x for + # Kimi-K2.5), but `warmup_compile_prefill` is idempotent: each config + # is only JIT'd once and cached in a module-global dict. + # tokenspeed_mla requires --kv-cache-dtype fp8_e4m3, so tokenspeed's + # FP8 prefill path (deepseek_v3.py:946 `use_fp8_prefill`) is always + # on and feeds fp8_e4m3fn q/k/v to the kernel — bf16 is unreachable + # for this backend. + d_qk = self.qk_nope_head_dim + self.qk_rope_head_dim + warmup_compile_prefill( + q_dtype=torch.float8_e4m3fn, + d_qk=d_qk, + d_v=self.v_head_dim, + enable_pdl=pdl_enabled(), + ) + + # Validate page_size + if self.page_size not in (32, 64): + raise ValueError( + f"tokenspeed_mla backend requires page_size 32 or 64, got {self.page_size}" + ) + + # tokenspeed_mla's CuTe DSL kernel only supports fp8_e4m3 KV cache; check + # at startup so misconfiguration surfaces here, not in the first forward. + kv_cache_dtype = global_server_args_dict.get("kv_cache_dtype", "auto") + if kv_cache_dtype != "fp8_e4m3": + raise NotImplementedError( + f"tokenspeed_mla backend requires --kv-cache-dtype fp8_e4m3, " + f"got {kv_cache_dtype!r}." + ) + + self.num_local_heads = num_heads_per_tp + + # Metadata + self.forward_decode_metadata: CuteDSLMLADecodeMetadata | None = None + self.forward_prefill_metadata: CuteDSLMLAPrefillMetadata | None = None + self.decode_cuda_graph_metadata: dict[int, CuteDSLMLADecodeMetadata] = {} + self.decode_cuda_graph_kv_indices = None + self.chunked_prefill_metadata: TRTLLMMLAChunkedPrefillMetadata | None = None + + def _calc_padded_blocks(self, max_seq_len: int) -> int: + """Calculate block count padded to satisfy the fused-kernel constraint.""" + blocks = triton.cdiv(max_seq_len, self.page_size) + constraint = TRTLLM_BLOCK_CONSTRAINT // self.page_size + if blocks % constraint != 0: + blocks = triton.cdiv(blocks, constraint) * constraint + return blocks + + def _create_block_kv_indices( + self, + batch_size: int, + max_blocks: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + req_to_page: torch.Tensor, + block_kv_indices: torch.Tensor | None = None, + ) -> torch.Tensor: + """Build page-table from req_to_page using vectorized tensor indexing.""" + if block_kv_indices is None: + block_kv_indices = torch.zeros( + (batch_size, max_blocks), dtype=torch.int32, device=self.device + ) + + copy_len = min(max_blocks, req_to_page.shape[1]) + + block_kv_indices[:batch_size, :copy_len] = req_to_page[ + req_pool_indices[:batch_size], :copy_len + ] + + return block_kv_indices + + # ---- Metadata initialization ---- + + def init_forward_metadata( + self, + bs: int, + num_extends: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + forward_mode: ForwardMode, + req_to_page: torch.Tensor, + seq_lens_cpu: torch.Tensor | None = None, + spec_info=None, + **kwargs, + ): + if forward_mode.is_extend_or_mixed(): + self._init_prefill_metadata( + seq_lens[:num_extends], + req_pool_indices=req_pool_indices[:num_extends], + req_to_page=req_to_page, + extend_prefix_lens=kwargs.pop("extend_prefix_lens"), + extend_prefix_lens_cpu=kwargs.pop("extend_prefix_lens_cpu"), + extend_seq_lens=kwargs.pop("extend_seq_lens"), + extend_seq_lens_cpu=kwargs.pop("extend_seq_lens_cpu"), + ) + # Drafter steps 1..N are pure DECODE on full bs regardless of target + # mode, so under is_draft we also fill decode_metadata under EXTEND + # so the multi-step loop has metadata. The wrapper pre-writes + # draft_seq_lens before calling here so `seq_lens` aliases the + # drafter's live buffer. + if ( + forward_mode.is_decode() + or forward_mode.is_mixed() + or (forward_mode.is_extend() and self.is_draft) + ): + self._init_decode_metadata( + bs, + num_extends, + req_pool_indices, + seq_lens, + req_to_page, + ) + + @contextmanager + def override_num_extends(self, num_extends: int): + assert self.forward_decode_metadata is not None + prev = self.forward_decode_metadata.num_extends + self.forward_decode_metadata.num_extends = num_extends + try: + yield + finally: + self.forward_decode_metadata.num_extends = prev + + def _init_decode_metadata( + self, + bs: int, + num_extends: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + req_to_page: torch.Tensor, + ): + max_blocks = self._calc_padded_blocks(self.max_context_len) + block_kv_indices = self._create_block_kv_indices( + bs, max_blocks, req_pool_indices, seq_lens, req_to_page + ) + + self.forward_decode_metadata = CuteDSLMLADecodeMetadata( + block_kv_indices=block_kv_indices, + max_seq_len_k=self.max_context_len, + seq_lens_k=seq_lens, + num_extends=num_extends, + ) + + def _init_prefill_metadata( + self, + seq_lens: torch.Tensor, + req_pool_indices: torch.Tensor | None = None, + req_to_page: torch.Tensor | None = None, + extend_prefix_lens: torch.Tensor | None = None, + extend_prefix_lens_cpu: torch.Tensor | None = None, + extend_seq_lens: torch.Tensor | None = None, + extend_seq_lens_cpu: torch.Tensor | None = None, + ): + # Worst-case bound to avoid GPU->CPU sync from seq_lens.max().item(). + # TODO: track a loose CPU upper bound (advance by chunked_prefill_size / + # accept_lengths.max(); correct when accurate values land) for tighter + # kernel-grid sizing without syncing. + max_seq_len = self.max_context_len + cum_seq_lens = torch.zeros( + len(seq_lens) + 1, dtype=torch.int32, device=seq_lens.device + ) + torch.cumsum(seq_lens, dim=0, out=cum_seq_lens[1:]) + + assert ( + seq_lens.dtype == torch.int32 + ), f"seq_lens must be int32, got {seq_lens.dtype}" + self.forward_prefill_metadata = CuteDSLMLAPrefillMetadata( + max_seq_len=max_seq_len, + cum_seq_lens=cum_seq_lens, + seq_lens=seq_lens, + ) + num_extends = extend_seq_lens.shape[0] + cum_extend_seq_lens = torch.zeros( + num_extends + 1, device=self.device, dtype=torch.int32 + ) + torch.cumsum(extend_seq_lens, dim=0, out=cum_extend_seq_lens[1:]) + max_extend_seq_len = extend_seq_lens_cpu.max().item() + ( + chunked_loop_num, + chunk_kv_indices_list, + chunked_seq_len, + cu_chunked_seq_len, + max_chunk_len_per_loop, + ) = build_chunked_prefill_metadata_arrays( + extend_prefix_lens, + extend_prefix_lens_cpu, + req_to_page, + req_pool_indices, + self.page_size, + ) + self.chunked_prefill_metadata = TRTLLMMLAChunkedPrefillMetadata( + extend_prefix_lens=extend_prefix_lens, + extend_prefix_lens_cpu=extend_prefix_lens_cpu, + extend_seq_lens=extend_seq_lens, + extend_seq_lens_cpu=extend_seq_lens_cpu, + req_pool_indices=req_pool_indices, + cum_extend_seq_lens=cum_extend_seq_lens, + max_extend_seq_len=max_extend_seq_len, + chunked_loop_num=chunked_loop_num, + chunk_kv_indices_list=chunk_kv_indices_list, + chunked_seq_len=chunked_seq_len, + cu_chunked_seq_len=cu_chunked_seq_len, + max_chunk_len_per_loop=max_chunk_len_per_loop, + ) + + # ---- CUDA Graph ---- + + def init_cuda_graph_state(self, max_bs: int, seq_lens_buf: torch.Tensor): + assert ( + seq_lens_buf.dtype == torch.int32 + and seq_lens_buf.dim() == 1 + and seq_lens_buf.shape[0] >= max_bs + ), ( + f"seq_lens_buf must be int32 with shape[0] >= {max_bs}, " + f"got {seq_lens_buf.dtype} {tuple(seq_lens_buf.shape)}" + ) + # Alias controller's seq_lens_buf — backend never mutates it. + self.cuda_graph_seq_lens_buf = seq_lens_buf + max_blocks = self._calc_padded_blocks(self.max_context_len) + self.decode_cuda_graph_kv_indices = torch.zeros( + (max_bs, max_blocks), dtype=torch.int32, device=self.device + ) + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + forward_mode: ForwardMode, + ): + if forward_mode.is_extend_or_mixed(): + raise NotImplementedError( + f"tokenspeed_mla CUDA graph capture not supported for {forward_mode}" + ) + + max_blocks = self._calc_padded_blocks(self.max_context_len) + block_kv_indices = self.decode_cuda_graph_kv_indices[:bs, :max_blocks] + + metadata = CuteDSLMLADecodeMetadata( + block_kv_indices=block_kv_indices, + max_seq_len_k=self.max_context_len, + seq_lens_k=self.cuda_graph_seq_lens_buf[:bs], + num_extends=0, + ) + 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, + forward_mode: ForwardMode = None, + req_to_page: torch.Tensor = None, + **kwargs, + ): + if forward_mode is not None and forward_mode.is_extend_or_mixed(): + raise NotImplementedError( + f"tokenspeed_mla CUDA graph replay not supported for {forward_mode}" + ) + + metadata = self.decode_cuda_graph_metadata[bs] + + # seq_lens_k aliases seq_lens_buf; only block indices need refresh. + # When the buffer is aliased to a peer backend (e.g. drafter aliasing + # the target's kv_indices), the peer's replay has already populated it + # with identical content. + if req_to_page is not None and not self._block_table_aliased: + self._create_block_kv_indices( + bs, + metadata.block_kv_indices.shape[1], + req_pool_indices[:bs], + seq_lens[:bs], + req_to_page, + metadata.block_kv_indices, + ) + + self.forward_decode_metadata = metadata + + def get_cuda_graph_seq_len_fill_value(self): + return 1 + + # ---- Forward: Decode ---- + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: PagedAttention, + out_cache_loc: torch.Tensor, + token_to_kv_pool, + bs: int, + save_kv_cache: bool = True, + **kwargs, + ) -> torch.Tensor: + # q is whole Q [T, H, head_dim]; k is whole latent [T, 1, head_dim]. + if save_kv_cache: + assert k is not None + token_to_kv_pool.set_mla_kv_buffer( + layer, + out_cache_loc, + k[..., : self.kv_lora_rank], + k[..., self.kv_lora_rank :], + ) + + metadata = self.forward_decode_metadata + num_extends = metadata.num_extends + q_len_per_req = q.shape[0] // bs + query = q.view(bs, q_len_per_req, layer.tp_q_head_num, layer.head_dim) + + softmax_scale = layer.scaling + if self.data_type == torch.float8_e4m3fn: + query = query.to(self.data_type) + k_scale = ( + layer.k_scale_float + if getattr(layer, "k_scale_float", None) is not None + else 1.0 + ) + softmax_scale = k_scale * layer.scaling + + # Prepare KV cache: [num_pages, page_size, kv_cache_dim] (3D for CuteDSL) + k_cache = token_to_kv_pool.get_key_buffer(layer.layer_id) + if self.data_type != k_cache.dtype: + k_cache = k_cache.to(self.data_type) + kv_cache = k_cache.view(-1, self.page_size, self.kv_cache_dim) + + if not CuteDSLMLABackend._logged_decode: + logger.info( + "CuteDSL MLA decode kernel invoked (tokenspeed_mla_decode, query_dtype=%s, kv_dtype=%s)", + query.dtype, + kv_cache.dtype, + ) + CuteDSLMLABackend._logged_decode = True + + self.cutedsl_workspace = get_cutedsl_workspace_buffer( + query.device, layer.tp_q_head_num, self.kv_lora_rank, query.shape[1] + ) + + raw_out = tokenspeed_mla_decode( + query=query, + kv_cache=kv_cache, + workspace_buffer=self.cutedsl_workspace, + kv_lora_rank=self.kv_lora_rank, + qk_rope_head_dim=self.qk_rope_head_dim, + block_tables=metadata.block_kv_indices[num_extends:], + seq_lens=metadata.seq_lens_k[num_extends:], + max_seq_len=metadata.max_seq_len_k, + softmax_scale=softmax_scale, + enable_pdl=pdl_enabled(), + ) + + return raw_out.view(-1, layer.tp_q_head_num * layer.v_head_dim) + + # ---- Forward: Extend/Prefill ---- + + def forward_extend_chunked( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + scaling, + logits_soft_cap, + *, + cum_seq_lens_q, + cum_seq_lens_kv, + max_q_len, + max_kv_len, + seq_lens, + batch_size, + causal, + out: torch.Tensor | None = None, + ): + if causal: + step_counter = getattr(self, "step_counter", None) + if step_counter is not None: + step_counter.record_cache() + + head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + q = q.reshape(-1, self.num_local_heads, head_dim) + k = k.reshape(-1, self.num_local_heads, head_dim) + v = v.reshape(-1, self.num_local_heads, self.v_head_dim) + + # CuteDSL FMHA MLA: if Q is FP8, ensure K/V match. `.to()` is a no-op + # when the source dtype already matches. + if q.dtype == torch.float8_e4m3fn: + k = k.to(torch.float8_e4m3fn) + v = v.to(torch.float8_e4m3fn) + + if not CuteDSLMLABackend._logged_prefill: + logger.info("CuteDSL MLA prefill kernel invoked (tokenspeed_mla_prefill)") + CuteDSLMLABackend._logged_prefill = True + + result = tokenspeed_mla_prefill( + query=q, + key=k, + value=v, + seq_lens=seq_lens, + cum_seq_lens=cum_seq_lens_kv, + max_seq_len=max_kv_len, + batch_size=batch_size, + softmax_scale=scaling, + is_causal=causal, + return_lse=True, + cum_seq_lens_q=cum_seq_lens_q, + max_seq_len_q=max_q_len, + enable_pdl=pdl_enabled(), + out=out, + ) + + if isinstance(result, tuple): + out, lse = result[0], result[1] + else: + out, lse = result, None + + if out.dtype != self.q_data_type: + out = out.to(self.q_data_type) + + return out, lse + + +register_backend("tokenspeed_mla", {AttentionArch.MLA}, CuteDSLMLABackend) diff --git a/python/tokenspeed/runtime/layers/attention/backends/trtllm.py b/python/tokenspeed/runtime/layers/attention/backends/trtllm.py new file mode 100644 index 0000000..6d33110 --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/backends/trtllm.py @@ -0,0 +1,965 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +""" +MHA attention backend for TokenSpeed scheduling. +Uses fused kernels optimized for SM100 (Blackwell). +Supports sliding window, attention sinks, and FP8 KV cache. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch +from tokenspeed_kernel.ops.attention.flashinfer import ( + trtllm_batch_context_with_kv_cache, + trtllm_batch_decode_with_kv_cache, +) +from tokenspeed_kernel.ops.kvcache.triton import ( + fused_fp8_set_kv_buffer, + gather_page_table_with_padding, +) + +from tokenspeed.runtime.configs.model_config import AttentionArch +from tokenspeed.runtime.execution.forward_batch_info import ForwardMode +from tokenspeed.runtime.layers.attention.backends.base import AttentionBackend +from tokenspeed.runtime.layers.attention.backends.flat_groups import ( + FlatCacheGroupsMixin, +) +from tokenspeed.runtime.layers.attention.configs.mha import MHAConfig +from tokenspeed.runtime.layers.attention.registry import register_backend +from tokenspeed.runtime.layers.common import fp8_cast_contiguous +from tokenspeed.runtime.utils import get_colorful_logger + +if TYPE_CHECKING: + from tokenspeed.runtime.layers.paged_attention import PagedAttention + +logger = get_colorful_logger(__name__) + +# Workspace buffer shared across all trtllm_mha wrappers. +_global_workspace_buffer: torch.Tensor | None = None +TRTLLM_MHA_WORKSPACE = 512 * 1024 * 1024 + + +def canonicalize_stride(tensor: torch.Tensor) -> torch.Tensor: + """Adjust degenerate strides for a tensor, make it canonical. + + When a dimension has size=1, PyTorch may use the same stride as the next dim. + This causes TMA desc validation failures in the trtllm_mha backend. + See: https://github.com/flashinfer-ai/flashinfer/issues/2232 + """ + sizes = tensor.size() + strides = tensor.stride() + ndim = tensor.dim() + + need_fix = any( + sizes[i] == 1 and strides[i] == strides[i + 1] for i in range(ndim - 1) + ) + + if not need_fix: + return tensor + + new_strides = [0] * ndim + new_strides[-1] = 1 + for i in range(ndim - 2, -1, -1): + new_strides[i] = new_strides[i + 1] * sizes[i + 1] + + return tensor.as_strided(sizes, new_strides) + + +@dataclass +class TRTLLMMHAMetadata: + cache_seqlens_int32: torch.Tensor = None + max_seq_len_q: int = 1 + max_seq_len_k: int = 0 + cu_seqlens_q: torch.Tensor = None + cu_seqlens_k: torch.Tensor = None + # page_table is None on the flat path (per-group page_tables route reads). + page_table: torch.Tensor = None + # Flat per-group tables/write-locs, keyed by group id (see flat_groups). + page_tables: dict[str, torch.Tensor] | None = None + out_cache_locs: dict[str, torch.Tensor] | None = None + + +class TRTLLMMHAAttnBackend(FlatCacheGroupsMixin, AttentionBackend): + """trtllm_mha attention backend optimized for SM100 (Blackwell).""" + + # Per-group flat tables: reads and writes route by layer.group_id, so + # slab-aliased pools (e.g. gpt-oss sliding+full pairing) are safe. + uses_flat_cache_groups: bool = True + # Graph-buffer column tails pad with the zero-init dummy page, matching + # the radix replay contract (gather_page_table_with_padding dummy_slot=0). + flat_table_tail_pad: int = 0 + + def support_kv_cache_prewrite( + self, forward_mode: ForwardMode | None = None + ) -> bool: + return True + + def _prewrite_metadata(self, forward_mode): + # Prewrite fires on extend too (unlike MHA): route it to the slot + # init_forward_metadata filled for this forward. Target verify is + # DECODE mode but its multi-token metadata lives in the prefill slot. + if forward_mode is not None and forward_mode.is_extend_or_mixed(): + return self.forward_prefill_metadata + if self.spec_num_tokens > 1 and not self.is_draft: + return self.forward_prefill_metadata + return self.forward_decode_metadata + + @property + def sinks_dtype(self) -> torch.dtype: + return torch.float32 + + def __init__(self, config: MHAConfig): + super().__init__(config) + + self.page_size = config.page_size + self.max_context_len = config.context_len + self.kv_cache_dtype = config.kv_cache_dtype + max_bs = config.max_bs + + # Shared workspace buffer (allocated once per process). + global _global_workspace_buffer + if _global_workspace_buffer is None: + _global_workspace_buffer = torch.zeros( + TRTLLM_MHA_WORKSPACE, + dtype=torch.uint8, + device=config.device, + ) + self.workspace_buffer = _global_workspace_buffer + + # Max pages per request. + self.max_num_pages = (config.context_len + self.page_size - 1) // self.page_size + + # Persistent buffers for page table construction. + self.page_table_buf = torch.zeros( + (max_bs, self.max_num_pages), + dtype=torch.int32, + device=config.device, + ) + self.cache_seqlens_buf = torch.zeros( + (max_bs,), dtype=torch.int32, device=config.device + ) + # KV seqlens clamped to >= spec_num_tokens for the MTP verify path. + # Padded decode rows have seq_len=1 (InputBuffer); with q_len=spec_num_tokens + # they'd hit an empty causal span and the kernel returns NaN. Mirrors mha.py. + self.spec_cache_seqlens_buf = torch.zeros( + (max_bs,), dtype=torch.int32, device=config.device + ) + self.cu_seqlens_q_buf = torch.zeros( + (max_bs + 1,), dtype=torch.int32, device=config.device + ) + self.cu_seqlens_k_buf = torch.zeros( + (max_bs + 1,), dtype=torch.int32, device=config.device + ) + + # Separate slots for prefill-kernel vs decode-kernel forward paths. + # forward_extend reads prefill; forward_decode reads decode. + self.forward_prefill_metadata: TRTLLMMHAMetadata | None = None + self.forward_decode_metadata: TRTLLMMHAMetadata | None = None + + # CUDA graph state — per-slot dicts. + self.cuda_graph_prefill_metadata: dict[int, TRTLLMMHAMetadata] = {} + self.cuda_graph_decode_metadata: dict[int, TRTLLMMHAMetadata] = {} + + # DFLASH draft: the drafter predicts a whole block of spec_num_tokens + # per decode forward and needs non-causal (block-diffusion) attention. + # Instead of a non-causal mask, expand each request into spec_num_tokens + # single-query rows sharing the SAME block-end seq_len, so each row + # attends over the whole block. Mirrors the MHA draft_block_decode path; + # target verify and ordinary trtllm decode are untouched. + self.draft_block_decode = bool(getattr(config, "draft_block_decode", False)) + + # ------------------------------------------------------------------ + # Page table helpers + # ------------------------------------------------------------------ + + def _build_page_table( + self, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + bs: int, + req_to_page: torch.Tensor, + page_table_buf: torch.Tensor, + ) -> torch.Tensor: + """Build page table in [bs, max_pages] format from req_to_page. + + req_to_page is [req_pool_size+1, max_pages] containing page IDs. + """ + page_table_buf[:bs].copy_( + req_to_page[req_pool_indices[:bs], : self.max_num_pages] + ) + return page_table_buf[:bs] + + # ------------------------------------------------------------------ + # KV cache helpers + # ------------------------------------------------------------------ + + def _get_kv_cache_permuted(self, layer: PagedAttention, token_to_kv_pool): + """Get KV cache in [num_pages, num_kv_heads, page_size, head_dim] layout.""" + k_cache, v_cache = 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) + + if layer.tp_k_head_num == 1: + k_cache = canonicalize_stride(k_cache) + if layer.tp_v_head_num == 1: + v_cache = canonicalize_stride(v_cache) + + return k_cache, v_cache + + def _compute_scales(self, layer: PagedAttention): + """Compute bmm1/bmm2 scales for the fused kernel.""" + 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 + return bmm1_scale, bmm2_scale + + def _should_use_fused_fp8_path(self, save_kv_cache: bool, k) -> bool: + return ( + save_kv_cache + and k is not None + and self.kv_cache_dtype == torch.float8_e4m3fn + ) + + # ------------------------------------------------------------------ + # Forward + # ------------------------------------------------------------------ + + def _save_kv_and_prepare_q( + self, q, k, v, layer, out_cache_loc, token_to_kv_pool, save_kv_cache + ): + if self._should_use_fused_fp8_path(save_kv_cache, k): + k_cache, v_cache = token_to_kv_pool.get_kv_buffer(layer.layer_id) + fused_fp8_set_kv_buffer( + k=k.view(-1, layer.tp_k_head_num, layer.head_dim), + v=v.view(-1, layer.tp_k_head_num, layer.head_dim), + k_cache=k_cache, + v_cache=v_cache, + cache_loc=out_cache_loc, + k_scale=layer.k_scale, + v_scale=layer.v_scale, + page_size=self.page_size, + ) + elif save_kv_cache and k is not None: + token_to_kv_pool.set_kv_buffer( + layer, out_cache_loc, k, v, layer.k_scale, layer.v_scale + ) + + if self.kv_cache_dtype == torch.float8_e4m3fn: + q = fp8_cast_contiguous(q) + else: + q = q.contiguous() + + return q.view(-1, layer.tp_q_head_num, layer.head_dim) + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: PagedAttention, + out_cache_loc: torch.Tensor, + token_to_kv_pool, + bs: int, + save_kv_cache: bool = True, + **kwargs, + ) -> torch.Tensor: + if self.draft_block_decode and self.spec_num_tokens > 1: + # DFLASH draft block: metadata is expanded to bs*spec_num_tokens + # single-query rows, so use the decode slot directly. Inferring + # q_len_per_req from q.shape[0]//bs would be spec_num_tokens and + # wrongly pick the prefill slot. + metadata = self.forward_decode_metadata + else: + # Multi-token decode (q_len > 1) reads the prefill slot's + # uniform-stride metadata; plain decode reads the single-token slot. + q_len_per_req = q.shape[0] // bs if bs > 0 else 1 + metadata = ( + self.forward_prefill_metadata + if q_len_per_req > 1 + else self.forward_decode_metadata + ) + + out_cache_loc = self._select_out_cache_loc( + layer, metadata, out_cache_loc, prefer_caller=self.is_draft + ) + q = self._save_kv_and_prepare_q( + q, k, v, layer, out_cache_loc, token_to_kv_pool, save_kv_cache + ) + k_cache, v_cache = self._get_kv_cache_permuted(layer, token_to_kv_pool) + bmm1_scale, bmm2_scale = self._compute_scales(layer) + + attention_sink = kwargs.get("sinks", None) + if attention_sink is not None: + attention_sink = attention_sink.float() + + o = trtllm_batch_decode_with_kv_cache( + query=q, + kv_cache=(k_cache, v_cache), + workspace_buffer=self.workspace_buffer, + block_tables=self._select_page_table(layer, metadata), + seq_lens=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, + sinks=attention_sink, + out_dtype=self.dtype, + q_len_per_req=metadata.max_seq_len_q, + ) + 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: PagedAttention, + out_cache_loc: torch.Tensor, + token_to_kv_pool, + bs: int, + save_kv_cache: bool = True, + **kwargs, + ) -> torch.Tensor: + metadata = self.forward_prefill_metadata + out_cache_loc = self._select_out_cache_loc(layer, metadata, out_cache_loc) + q = self._save_kv_and_prepare_q( + q, k, v, layer, out_cache_loc, token_to_kv_pool, save_kv_cache + ) + k_cache, v_cache = self._get_kv_cache_permuted(layer, token_to_kv_pool) + bmm1_scale, bmm2_scale = self._compute_scales(layer) + + attention_sink = kwargs.get("sinks", None) + if attention_sink is not None: + attention_sink = attention_sink.float() + + o = trtllm_batch_context_with_kv_cache( + query=q, + kv_cache=(k_cache, v_cache), + workspace_buffer=self.workspace_buffer, + block_tables=self._select_page_table(layer, metadata), + seq_lens=metadata.cache_seqlens_int32, + max_q_len=metadata.max_seq_len_q, + max_kv_len=self.max_context_len, + bmm1_scale=bmm1_scale, + bmm2_scale=bmm2_scale, + batch_size=metadata.cu_seqlens_q.shape[0] - 1, + cum_seq_lens_q=metadata.cu_seqlens_q, + cum_seq_lens_kv=metadata.cu_seqlens_k, + window_left=layer.sliding_window_size, + sinks=attention_sink, + out_dtype=self.dtype, + ) + return o.view(-1, layer.tp_q_head_num * layer.head_dim) + + # ------------------------------------------------------------------ + # Metadata initialisation + # ------------------------------------------------------------------ + + def init_forward_metadata( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + forward_mode: ForwardMode, + req_to_page: torch.Tensor, + extend_with_prefix: bool = False, + extend_prefix_lens: torch.Tensor | None = None, + extend_prefix_lens_cpu: torch.Tensor | None = None, + extend_seq_lens_cpu: torch.Tensor | None = None, + spec_info=None, + use_cuda_graph: bool = False, + flat_block_tables: dict[str, torch.Tensor] | None = None, + **kwargs, + ): + flat_page_tables = self._shed_state_groups(flat_block_tables) + flat_out_cache_locs = None + if flat_page_tables: + # Verify keeps [bs]-row tables; only DFLASH expands rows. TODO(flat+dflash). + assert not ( + self.draft_block_decode and self.spec_num_tokens > 1 + ), "flat cache groups are unsupported with DFLASH block decode" + if forward_mode.is_extend_or_mixed(): + assert extend_prefix_lens_cpu is not None + assert extend_seq_lens_cpu is not None + flat_out_cache_locs = self._compute_flat_extend_out_cache_locs( + flat_page_tables, + extend_prefix_lens_cpu[:bs], + extend_seq_lens_cpu[:bs], + self.page_size, + ) + else: + flat_out_cache_locs = self._compute_flat_decode_out_cache_locs( + flat_page_tables, + seq_lens[:bs], + self.page_size, + self._flat_verify_tokens(), + ) + self._maybe_check_flat_write_locs( + flat_page_tables, flat_out_cache_locs, self.page_size + ) + + if forward_mode.is_extend_or_mixed(): + self._init_extend_metadata( + bs, + req_pool_indices, + seq_lens, + req_to_page, + extend_with_prefix=extend_with_prefix, + extend_prefix_lens=extend_prefix_lens, + extend_prefix_lens_cpu=extend_prefix_lens_cpu, + extend_seq_lens_cpu=extend_seq_lens_cpu, + flat_page_tables=flat_page_tables, + flat_out_cache_locs=flat_out_cache_locs, + ) + # Drafter: also fill decode_metadata so step 1+ multi-step has + # metadata under EXTEND/MIXED target. seq_lens is the drafter's + # live alias buffer (wrapper pre-writes before this call). + if self.is_draft: + self._init_decode_metadata(bs, req_pool_indices, seq_lens, req_to_page) + return + + if self.draft_block_decode and self.spec_num_tokens > 1: + # DFLASH draft block (eager): expand to spec_num_tokens single-query + # rows per request; seq_lens is the block-end length the drafter + # already wrote. + self._init_block_decode_metadata( + bs, req_pool_indices, seq_lens, req_to_page + ) + return + + if self.spec_num_tokens > 1: + self._init_multi_token_metadata( + bs, + self.spec_num_tokens, + req_pool_indices, + seq_lens, + req_to_page, + flat_page_tables=flat_page_tables, + flat_out_cache_locs=flat_out_cache_locs, + ) + if self.is_draft: + # Drafter's N-1 single-token steps after the first. + self._init_decode_metadata( + bs, + req_pool_indices, + seq_lens, + req_to_page, + flat_page_tables=flat_page_tables, + flat_out_cache_locs=flat_out_cache_locs, + ) + else: + self._init_decode_metadata( + bs, + req_pool_indices, + seq_lens, + req_to_page, + flat_page_tables=flat_page_tables, + flat_out_cache_locs=flat_out_cache_locs, + ) + + def _init_decode_metadata( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + req_to_page: torch.Tensor, + flat_page_tables: dict[str, torch.Tensor] | None = None, + flat_out_cache_locs: dict[str, torch.Tensor] | None = None, + ): + assert ( + seq_lens.dtype == torch.int32 + ), f"seq_lens must be int32, got {seq_lens.dtype}" + device = seq_lens.device + # Alias seq_lens (no copy, no mutation). cu_seqlens_k omitted: + # the decode kernel doesn't read it. On the flat path the per-group + # tables route every read; the radix single table would be dead work. + self.forward_decode_metadata = TRTLLMMHAMetadata( + cache_seqlens_int32=seq_lens[:bs], + max_seq_len_q=1, + max_seq_len_k=self.max_context_len, + cu_seqlens_q=torch.arange(0, bs + 1, dtype=torch.int32, device=device), + page_table=( + None + if flat_page_tables + else self._build_page_table( + req_pool_indices, seq_lens, bs, req_to_page, self.page_table_buf + ) + ), + page_tables=flat_page_tables, + out_cache_locs=flat_out_cache_locs, + ) + + def _replicate_block_page_table( + self, + out: torch.Tensor, + req_pool_indices: torch.Tensor, + bs: int, + req_to_page: torch.Tensor, + ) -> None: + """Replicate each request's page table to its spec_num_tokens block rows. + + ``out`` is the [bs*spec_num_tokens, max_num_pages] destination. All block + rows of a request share its pages (decode only reads KV), so a single + broadcast copy suffices. + """ + spec = self.spec_num_tokens + base_page_table = req_to_page[req_pool_indices[:bs], : self.max_num_pages] + out[: bs * spec, :].view(bs, spec, self.max_num_pages).copy_( + base_page_table[:, None, :] + ) + + def _init_block_decode_metadata( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + req_to_page: torch.Tensor, + ): + """Eager DFLASH draft-block metadata: spec_num_tokens single-query rows + per request, all carrying the block-end seq_len (prefix + spec_num_tokens) + so each query attends over the whole block. Allocates fresh buffers (the + cuda-graph path uses persistent ones), mirroring the MHA backend. + """ + assert ( + seq_lens.dtype == torch.int32 + ), f"seq_lens must be int32, got {seq_lens.dtype}" + spec = self.spec_num_tokens + device = seq_lens.device + expanded_bs = bs * spec + + page_table = torch.empty( + (expanded_bs, self.max_num_pages), dtype=torch.int32, device=device + ) + self._replicate_block_page_table(page_table, req_pool_indices, bs, req_to_page) + + # Clamp the block-end length so the decode never asks for more page-table + # columns than exist (prefix + spec_num_tokens can exceed max_context_len). + cache_seqlens = ( + seq_lens[:bs] + .clamp(spec, self.max_context_len) + .unsqueeze(1) + .expand(bs, spec) + .reshape(expanded_bs) + .contiguous() + ) + + self.forward_decode_metadata = TRTLLMMHAMetadata( + cache_seqlens_int32=cache_seqlens, + max_seq_len_q=1, + max_seq_len_k=self.max_context_len, + cu_seqlens_q=torch.arange( + 0, expanded_bs + 1, dtype=torch.int32, device=device + ), + page_table=page_table, + ) + + def _clamped_spec_seqlens( + self, seq_lens: torch.Tensor, bs: int, spec_num_tokens: int + ) -> torch.Tensor: + """Return KV seqlens clamped to >= spec_num_tokens for the MTP verify path. + + Writes into the persistent spec_cache_seqlens_buf (CUDA-graph safe) + to avoid NaN from empty causal spans on padded rows (seq_len=1). + """ + dst = self.spec_cache_seqlens_buf[:bs] + torch.clamp_min(seq_lens[:bs], spec_num_tokens, out=dst) + return dst + + def _flat_verify_tokens(self) -> int: + # Write locs per request on the flat path: N for target verify + # ([bs*N], token-major), 1 elsewhere (draft chains use caller locs). + return ( + self.spec_num_tokens + if self.spec_num_tokens > 1 and not self.is_draft + else 1 + ) + + def _init_multi_token_metadata( + self, + bs: int, + spec_num_tokens: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + req_to_page: torch.Tensor, + flat_page_tables: dict[str, torch.Tensor] | None = None, + flat_out_cache_locs: dict[str, torch.Tensor] | None = None, + ): + """Prefill-slot metadata for multi-token decode (uniform q_len per + request). Routes through the decode kernel via q_len_per_req; the + kernel doesn't read cu_seqlens_k.""" + assert ( + seq_lens.dtype == torch.int32 + ), f"seq_lens must be int32, got {seq_lens.dtype}" + device = seq_lens.device + self.forward_prefill_metadata = TRTLLMMHAMetadata( + cache_seqlens_int32=self._clamped_spec_seqlens( + seq_lens, bs, spec_num_tokens + ), + max_seq_len_q=spec_num_tokens, + max_seq_len_k=self.max_context_len, + cu_seqlens_q=torch.arange( + 0, + bs * spec_num_tokens + 1, + spec_num_tokens, + dtype=torch.int32, + device=device, + ), + page_table=( + None + if flat_page_tables + else self._build_page_table( + req_pool_indices, seq_lens, bs, req_to_page, self.page_table_buf + ) + ), + page_tables=flat_page_tables, + out_cache_locs=flat_out_cache_locs, + ) + + def _init_extend_metadata( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + req_to_page: torch.Tensor, + extend_with_prefix: bool = False, + extend_prefix_lens: torch.Tensor | None = None, + extend_prefix_lens_cpu=None, + extend_seq_lens_cpu=None, + flat_page_tables: dict[str, torch.Tensor] | None = None, + flat_out_cache_locs: dict[str, torch.Tensor] | None = None, + ): + """Populate prefill slot for regular EXTEND (ragged query).""" + assert ( + seq_lens.dtype == torch.int32 + ), f"seq_lens must be int32, got {seq_lens.dtype}" + assert ( + extend_seq_lens_cpu is not None + ), "trtllm extend requires extend_seq_lens_cpu (pinned-CPU mirror) to avoid GPU sync" + cache_seqlens_int32 = seq_lens[:bs] + cu_seqlens_k = torch.nn.functional.pad( + torch.cumsum(seq_lens, dim=0, dtype=torch.int32), (1, 0) + ) + # Flat path: per-group tables route every read (see _init_decode_metadata). + page_table = ( + None + if flat_page_tables + else self._build_page_table( + req_pool_indices, seq_lens, bs, req_to_page, self.page_table_buf + ) + ) + + # Read the max from the pinned-CPU mirror — avoids a per-iter + # GPU->CPU sync that would block the host on the previous step's + # forward and erase prefill/decode overlap. Both branches want + # max(new tokens per request); for a no-prefix extend that's + # seq_lens, for a prefix-cached extend it's seq_lens-prefix_lens — + # extend_seq_lens_cpu holds those new-token counts in either case. + max_seq_len_q = int(extend_seq_lens_cpu[:bs].max().item()) + + if extend_with_prefix and ( + (extend_prefix_lens_cpu is not None and any(extend_prefix_lens_cpu)) + or (extend_prefix_lens is not None and any(extend_prefix_lens.tolist())) + ): + if extend_prefix_lens is None: + raise RuntimeError( + "TRTLLMMHAAttnBackend requires extend_prefix_lens tensor " + "when extend_with_prefix is true." + ) + extend_seq_lens = seq_lens - extend_prefix_lens + cu_seqlens_q = torch.nn.functional.pad( + torch.cumsum(extend_seq_lens, dim=0, dtype=torch.int32), (1, 0) + ) + else: + cu_seqlens_q = cu_seqlens_k + + self.forward_prefill_metadata = TRTLLMMHAMetadata( + cache_seqlens_int32=cache_seqlens_int32, + max_seq_len_q=max_seq_len_q, + max_seq_len_k=self.max_context_len, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + page_table=page_table, + page_tables=flat_page_tables, + out_cache_locs=flat_out_cache_locs, + ) + + # ------------------------------------------------------------------ + # CUDA graph support + # ------------------------------------------------------------------ + + def init_cuda_graph_state( + self, + max_bs: int, + seq_lens_buf: torch.Tensor, + paged_cache_group_specs: Sequence = (), + **kwargs, + ): + assert ( + seq_lens_buf.dtype == torch.int32 + and seq_lens_buf.dim() == 1 + and seq_lens_buf.shape[0] >= max_bs + ), ( + f"seq_lens_buf must be int32 with shape[0] >= {max_bs}, " + f"got {seq_lens_buf.dtype} {tuple(seq_lens_buf.shape)}" + ) + self.cuda_graph_prefill_metadata = {} + self.cuda_graph_decode_metadata = {} + # Flat per-group persistent buffers + state-group shed; before the + # DFLASH early return (replay reads the dict for the stale guard). + self._learn_flat_state_groups(paged_cache_group_specs) + self._init_flat_graph_buffers(max_bs) + if self.draft_block_decode and self.spec_num_tokens > 1: + # DFLASH draft block: spec_num_tokens decode rows per request. Unlike + # the plain path, cache_seqlens is a dedicated buffer (NOT aliasing + # seq_lens_buf): it is filled in-graph by fill_block_decode_seq_lens. + self.cuda_graph_page_table = torch.zeros( + (max_bs * self.spec_num_tokens, self.max_num_pages), + dtype=torch.int32, + device=self.device, + ) + self.cuda_graph_cache_seqlens = torch.full( + (max_bs * self.spec_num_tokens,), + self.spec_num_tokens, + dtype=torch.int32, + device=self.device, + ) + return + # Alias controller's seq_lens_buf — backend never mutates it. + self.cuda_graph_page_table = torch.zeros( + (max_bs, self.max_num_pages), dtype=torch.int32, device=self.device + ) + self.cuda_graph_cache_seqlens = seq_lens_buf + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + forward_mode: ForwardMode, + flat_cache_group_ids: tuple[str, ...] = (), + **kwargs, + ): + if forward_mode.is_extend_or_mixed(): + raise NotImplementedError( + f"trtllm CUDA graph capture not supported for {forward_mode}" + ) + + # Real tables only arrive at replay: capture lazily allocates + # persistent per-group buffers and records metadata views into them. + if flat_cache_group_ids: + # Verify keeps [bs]-row tables + [bs*N] loc views. TODO(flat+dflash). + assert not ( + self.draft_block_decode and self.spec_num_tokens > 1 + ), "flat_cache_group_ids is unsupported with DFLASH block decode" + page_tables, out_cache_locs = self._flat_capture_group_views( + bs, flat_cache_group_ids, tokens_per_req=self._flat_verify_tokens() + ) + + if self.draft_block_decode and self.spec_num_tokens > 1: + self._init_block_decode_metadata_capture(bs) + return + + if self.spec_num_tokens > 1: + self._init_multi_token_metadata_capture( + bs, self.spec_num_tokens, page_tables, out_cache_locs + ) + if self.is_draft: + self._init_decode_metadata_capture( + bs, seq_lens, page_tables, out_cache_locs + ) + else: + self._init_decode_metadata_capture( + bs, seq_lens, page_tables, out_cache_locs + ) + + def _init_block_decode_metadata_capture(self, bs: int): + """DFLASH draft block (cuda-graph capture): spec_num_tokens single-query + rows per request over the persistent expanded buffers. seq_lens are + filled in-graph by fill_block_decode_seq_lens; seed a safe baseline here + so the capture run stays in range before that op records.""" + expanded_bs = bs * self.spec_num_tokens + self.cuda_graph_cache_seqlens[:expanded_bs].fill_(self.spec_num_tokens) + metadata = TRTLLMMHAMetadata( + cache_seqlens_int32=self.cuda_graph_cache_seqlens[:expanded_bs], + max_seq_len_q=1, + max_seq_len_k=self.max_context_len, + cu_seqlens_q=torch.arange( + 0, expanded_bs + 1, dtype=torch.int32, device=self.device + ), + page_table=self.cuda_graph_page_table[:expanded_bs, :], + ) + self.cuda_graph_decode_metadata[bs] = metadata + self.forward_decode_metadata = metadata + + def _init_decode_metadata_capture( + self, + bs: int, + seq_lens: torch.Tensor, + page_tables: dict[str, torch.Tensor] | None = None, + out_cache_locs: dict[str, torch.Tensor] | None = None, + ): + # cache_seqlens aliases seq_lens_buf (set in init_cuda_graph_state). + # Flat captures route reads through the per-group buffer views and + # replay never fills the radix single table, so record page_table=None + # instead of a slice of the never-filled zero buffer. + metadata = TRTLLMMHAMetadata( + cache_seqlens_int32=self.cuda_graph_cache_seqlens[:bs], + max_seq_len_q=1, + max_seq_len_k=self.max_context_len, + cu_seqlens_q=torch.arange(0, bs + 1, dtype=torch.int32, device=self.device), + page_table=( + None if page_tables is not None else self.cuda_graph_page_table[:bs, :] + ), + page_tables=page_tables, + out_cache_locs=out_cache_locs, + ) + self.cuda_graph_decode_metadata[bs] = metadata + self.forward_decode_metadata = metadata + + def _init_multi_token_metadata_capture( + self, + bs: int, + spec_num_tokens: int, + page_tables: dict[str, torch.Tensor] | None = None, + out_cache_locs: dict[str, torch.Tensor] | None = None, + ): + # Multi-token decode: seed spec_cache_seqlens_buf (clamped to >= + # spec_num_tokens) at capture so padded rows (seq_len=1) avoid NaN. + # The replay path refreshes it each step. + cache_seqlens = self._clamped_spec_seqlens( + self.cuda_graph_cache_seqlens, bs, spec_num_tokens + ) + metadata = TRTLLMMHAMetadata( + cache_seqlens_int32=cache_seqlens, + max_seq_len_q=spec_num_tokens, + max_seq_len_k=self.max_context_len, + cu_seqlens_q=torch.arange( + 0, + bs * spec_num_tokens + 1, + spec_num_tokens, + dtype=torch.int32, + device=self.device, + ), + page_table=( + None if page_tables is not None else self.cuda_graph_page_table[:bs, :] + ), + page_tables=page_tables, + out_cache_locs=out_cache_locs, + ) + self.cuda_graph_prefill_metadata[bs] = metadata + self.forward_prefill_metadata = metadata + + def init_forward_metadata_replay_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + forward_mode: ForwardMode, + req_to_page: torch.Tensor = None, + flat_block_tables: dict[str, torch.Tensor] | None = None, + **kwargs, + ): + if forward_mode.is_extend_or_mixed(): + raise NotImplementedError( + f"trtllm CUDA graph replay not supported for {forward_mode}" + ) + + # Fail loudly instead of replaying over stale/zero page tables. + self._flat_replay_stale_guard(bs, flat_block_tables) + + if self.draft_block_decode and self.spec_num_tokens > 1: + # DFLASH draft block: replicate the page table to each request's + # block rows. seq_lens are re-derived in-graph, so not touched here. + if req_to_page is not None: + self._replicate_block_page_table( + self.cuda_graph_page_table, req_pool_indices, bs, req_to_page + ) + if bs in self.cuda_graph_decode_metadata: + self.forward_decode_metadata = self.cuda_graph_decode_metadata[bs] + return + + # cache_seqlens aliases seq_lens_buf; only page tables need refresh. + # Flat captures read only the per-group buffers; the radix single + # table would be dead work there (and req_to_page is unpopulated on + # a flat scheduler build). + if not self.cuda_graph_flat_page_tables and req_to_page is not None: + gather_page_table_with_padding( + req_to_page=req_to_page, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + out=self.cuda_graph_page_table, + bs=bs, + max_num_pages=self.max_num_pages, + page_size=self.page_size, + dummy_slot=0, + ) + if flat_block_tables: + # cuda_graph_cache_seqlens aliases the controller's seq_lens_buf, + # filled by input prep BEFORE this call. + self._flat_replay_fill( + bs, + flat_block_tables, + self.cuda_graph_cache_seqlens, + tokens_per_req=self._flat_verify_tokens(), + ) + + # Refresh for both verify and draft: draft step 1 is multi-token + # and reads spec_cache_seqlens_buf; later single-token steps don't. + if self.spec_num_tokens > 1: + self._clamped_spec_seqlens(seq_lens, bs, self.spec_num_tokens) + + if bs in self.cuda_graph_prefill_metadata: + self.forward_prefill_metadata = self.cuda_graph_prefill_metadata[bs] + if bs in self.cuda_graph_decode_metadata: + self.forward_decode_metadata = self.cuda_graph_decode_metadata[bs] + + def fill_block_decode_seq_lens(self, bs: int, block_seq_lens: torch.Tensor) -> None: + """DFLASH: broadcast each request's block-end length to its + spec_num_tokens cuda-graph decode rows. + + Called by the drafter inside the captured graph so every replay + re-derives cache_seqlens from the live draft length. Mirrors the MHA + backend method of the same name. + + Args: + bs: Number of draft requests. + block_seq_lens: ``[bs]`` per-request block-end lengths. + """ + spec = self.spec_num_tokens + self.cuda_graph_cache_seqlens[: bs * spec].view(bs, spec).copy_( + block_seq_lens[:bs].clamp(spec, self.max_context_len).unsqueeze(1) + ) + + +register_backend("trtllm", {AttentionArch.MHA}, TRTLLMMHAAttnBackend) diff --git a/python/tokenspeed/runtime/layers/attention/backends/trtllm_mla.py b/python/tokenspeed/runtime/layers/attention/backends/trtllm_mla.py new file mode 100644 index 0000000..65e21af --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/backends/trtllm_mla.py @@ -0,0 +1,549 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +""" +MLA attention backend for TokenSpeed scheduling. + +Uses fused kernels optimized for SM100 (Blackwell) GPUs. +""" + +from __future__ import annotations + +import logging +from contextlib import contextmanager +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch +import triton +from tokenspeed_kernel.ops.attention.flashinfer import ( + trtllm_batch_decode_with_kv_cache_mla, + trtllm_ragged_attention_deepseek, +) + +from tokenspeed.runtime.configs.model_config import AttentionArch +from tokenspeed.runtime.execution.forward_batch_info import ForwardMode +from tokenspeed.runtime.layers.attention.backends.base import AttentionBackend +from tokenspeed.runtime.layers.attention.chunk import ( + build_chunked_prefill_metadata_arrays, +) +from tokenspeed.runtime.layers.attention.configs.mla import MLAConfig +from tokenspeed.runtime.layers.attention.registry import register_backend +from tokenspeed.runtime.utils.pdl import pdl_enabled + +if TYPE_CHECKING: + from tokenspeed.runtime.layers.paged_attention import PagedAttention + +logger = logging.getLogger(__name__) + +# Block constraint from flashinfer: block_num % (128 / page_size) == 0 +TRTLLM_BLOCK_CONSTRAINT = 128 + +# Shared workspace buffer for fused kernels (256 MB, zero-initialized). +# Zero-init is required for the kernel's internal semaphore mechanism. +_trtllm_workspace_buffer = None + + +def get_trtllm_workspace_buffer(device): + """Get or create the shared fused-kernel workspace buffer.""" + global _trtllm_workspace_buffer + if _trtllm_workspace_buffer is None: + _trtllm_workspace_buffer = torch.zeros( + 256 * 1024 * 1024, + dtype=torch.uint8, + device=device, + ) + return _trtllm_workspace_buffer + + +@dataclass +class TRTLLMMLAPrefillMetadata: + max_seq_len: int + cum_seq_lens: torch.Tensor + seq_lens: torch.Tensor + + +@dataclass +class TRTLLMMLAChunkedPrefillMetadata: + extend_prefix_lens: torch.Tensor + extend_prefix_lens_cpu: torch.Tensor + extend_seq_lens: torch.Tensor + extend_seq_lens_cpu: torch.Tensor + req_pool_indices: torch.Tensor + cum_extend_seq_lens: torch.Tensor # cumsum prefix-padded, sized num_extends+1 + max_extend_seq_len: int + # Per-prefix-chunk arrays for non-causal cross-attention (built once per + # iteration in _init_prefill_metadata, indexed by loop_idx in the model). + chunked_loop_num: int + chunk_kv_indices_list: list # List[torch.Tensor], one per loop_idx + chunked_seq_len: torch.Tensor # (chunked_loop_num, num_extends) int32 GPU + cu_chunked_seq_len: torch.Tensor # (chunked_loop_num, num_extends+1) int32 GPU + max_chunk_len_per_loop: list # List[int], one per loop_idx + # Per-request page table (req_to_page[req_pool_indices]). Populated only by + # the DSA backend for sparse-prefill top-k; plain MLA leaves it None. + block_tables: torch.Tensor | None = None + + +@dataclass +class TRTLLMMLADecodeMetadata: + num_extends: int = 0 + block_kv_indices: torch.Tensor | None = None + max_seq_len_k: int | None = None + seq_lens_k: torch.Tensor | None = None + + +class TRTLLMMLABackend(AttentionBackend): + """trtllm_mla attention backend using fused kernels.""" + + def __init__(self, config: MLAConfig): + super().__init__(config) + + self.max_context_len = config.context_len + self.page_size = config.page_size + + # MLA 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 = config.kv_cache_dim + self.scaling = config.scaling + self.data_type = config.kv_cache_dtype + self.q_data_type = config.dtype + + # Workspace zero-initialized for the fused kernel semaphore. + self.trtllm_workspace = get_trtllm_workspace_buffer(config.device) + + # Validate page_size + if self.page_size not in (32, 64): + raise ValueError( + f"trtllm_mla backend requires page_size 32 or 64, got {self.page_size}" + ) + + self.num_local_heads = config.num_attention_heads // config.attn_tp_size + + # Metadata + self.forward_decode_metadata: TRTLLMMLADecodeMetadata | None = None + self.forward_prefill_metadata: TRTLLMMLAPrefillMetadata | None = None + self.decode_cuda_graph_metadata: dict[int, TRTLLMMLADecodeMetadata] = {} + self.decode_cuda_graph_kv_indices = None + self.chunked_prefill_metadata: TRTLLMMLAChunkedPrefillMetadata | None = None + + def _calc_padded_blocks(self, max_seq_len: int) -> int: + """Calculate block count padded to satisfy the fused-kernel constraint.""" + blocks = triton.cdiv(max_seq_len, self.page_size) + constraint = TRTLLM_BLOCK_CONSTRAINT // self.page_size + if blocks % constraint != 0: + blocks = triton.cdiv(blocks, constraint) * constraint + return blocks + + def _create_block_kv_indices( + self, + batch_size: int, + max_blocks: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + req_to_page: torch.Tensor, + block_kv_indices: torch.Tensor | None = None, + ) -> torch.Tensor: + """Build page-table from req_to_page using vectorized tensor indexing.""" + if block_kv_indices is None: + block_kv_indices = torch.zeros( + (batch_size, max_blocks), dtype=torch.int32, device=self.device + ) + + copy_len = min(max_blocks, req_to_page.shape[1]) + + # Vectorized: gather all rows at once, no Python loop. + # Pages beyond actual seq_len are 0 (from req_to_page init); the kernel + # uses seq_lens to bound access so these padding entries are never read. + block_kv_indices[:batch_size, :copy_len] = req_to_page[ + req_pool_indices[:batch_size], :copy_len + ] + + return block_kv_indices + + # ---- Metadata initialization ---- + + def init_forward_metadata( + self, + bs: int, + num_extends: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + forward_mode: ForwardMode, + req_to_page: torch.Tensor, + spec_info=None, + **kwargs, + ): + if forward_mode.is_extend_or_mixed(): + self._init_prefill_metadata( + seq_lens[:num_extends], + req_pool_indices=req_pool_indices[:num_extends], + req_to_page=req_to_page, + extend_prefix_lens=kwargs.pop("extend_prefix_lens"), + extend_prefix_lens_cpu=kwargs.pop("extend_prefix_lens_cpu"), + extend_seq_lens=kwargs.pop("extend_seq_lens"), + extend_seq_lens_cpu=kwargs.pop("extend_seq_lens_cpu"), + ) + # Under is_draft, also fill decode_metadata under any forward_mode so + # the drafter's multi-step loop has metadata. Wrapper pre-writes + # draft_seq_lens before calling here, so `seq_lens` aliases the + # drafter's live buffer for step-1+ advances. + if ( + forward_mode.is_decode() + or forward_mode.is_mixed() + or (forward_mode.is_extend() and self.is_draft) + ): + self._init_decode_metadata( + bs, num_extends, req_pool_indices, seq_lens, req_to_page + ) + + @contextmanager + def override_num_extends(self, num_extends: int): + assert self.forward_decode_metadata is not None + prev = self.forward_decode_metadata.num_extends + self.forward_decode_metadata.num_extends = num_extends + try: + yield + finally: + self.forward_decode_metadata.num_extends = prev + + def _init_decode_metadata( + self, + bs: int, + num_extends: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + req_to_page: torch.Tensor, + ): + # For target_verify, the draft tokens have already been written to the KV + # cache. The seq_lens passed in should already reflect the full context. + # Use max_context_len to avoid GPU->CPU sync from seq_lens.max().item() + max_blocks = self._calc_padded_blocks(self.max_context_len) + + block_kv_indices = self._create_block_kv_indices( + bs, max_blocks, req_pool_indices, seq_lens, req_to_page + ) + + assert ( + seq_lens.dtype == torch.int32 + ), f"seq_lens must be int32, got {seq_lens.dtype}" + self.forward_decode_metadata = TRTLLMMLADecodeMetadata( + num_extends=num_extends, + block_kv_indices=block_kv_indices, + max_seq_len_k=self.max_context_len, + seq_lens_k=seq_lens, + ) + + def _init_prefill_metadata( + self, + seq_lens: torch.Tensor, + req_pool_indices: torch.Tensor | None = None, + req_to_page: torch.Tensor | None = None, + extend_prefix_lens: torch.Tensor | None = None, + extend_prefix_lens_cpu: torch.Tensor | None = None, + extend_seq_lens: torch.Tensor | None = None, + extend_seq_lens_cpu: torch.Tensor | None = None, + ): + max_seq_len = self.max_context_len + cum_seq_lens = torch.zeros( + len(seq_lens) + 1, dtype=torch.int32, device=seq_lens.device + ) + torch.cumsum(seq_lens, dim=0, out=cum_seq_lens[1:]) + + assert ( + seq_lens.dtype == torch.int32 + ), f"seq_lens must be int32, got {seq_lens.dtype}" + self.forward_prefill_metadata = TRTLLMMLAPrefillMetadata( + max_seq_len=max_seq_len, + cum_seq_lens=cum_seq_lens, + seq_lens=seq_lens, + ) + num_extends = extend_seq_lens.shape[0] + cum_extend_seq_lens = torch.zeros( + num_extends + 1, device=self.device, dtype=torch.int32 + ) + torch.cumsum(extend_seq_lens, dim=0, out=cum_extend_seq_lens[1:]) + max_extend_seq_len = extend_seq_lens_cpu.max().item() + ( + chunked_loop_num, + chunk_kv_indices_list, + chunked_seq_len, + cu_chunked_seq_len, + max_chunk_len_per_loop, + ) = build_chunked_prefill_metadata_arrays( + extend_prefix_lens, + extend_prefix_lens_cpu, + req_to_page, + req_pool_indices, + self.page_size, + ) + self.chunked_prefill_metadata = TRTLLMMLAChunkedPrefillMetadata( + extend_prefix_lens=extend_prefix_lens, + extend_prefix_lens_cpu=extend_prefix_lens_cpu, + extend_seq_lens=extend_seq_lens, + extend_seq_lens_cpu=extend_seq_lens_cpu, + req_pool_indices=req_pool_indices, + cum_extend_seq_lens=cum_extend_seq_lens, + max_extend_seq_len=max_extend_seq_len, + chunked_loop_num=chunked_loop_num, + chunk_kv_indices_list=chunk_kv_indices_list, + chunked_seq_len=chunked_seq_len, + cu_chunked_seq_len=cu_chunked_seq_len, + max_chunk_len_per_loop=max_chunk_len_per_loop, + ) + + # ---- CUDA Graph ---- + + def init_cuda_graph_state(self, max_bs: int, seq_lens_buf: torch.Tensor): + assert ( + seq_lens_buf.dtype == torch.int32 + and seq_lens_buf.dim() == 1 + and seq_lens_buf.shape[0] >= max_bs + ), ( + f"seq_lens_buf must be int32 with shape[0] >= {max_bs}, " + f"got {seq_lens_buf.dtype} {tuple(seq_lens_buf.shape)}" + ) + # Alias controller's seq_lens_buf — backend never mutates it. + self.cuda_graph_seq_lens_buf = seq_lens_buf + max_blocks = self._calc_padded_blocks(self.max_context_len) + self.decode_cuda_graph_kv_indices = torch.zeros( + (max_bs, max_blocks), dtype=torch.int32, device=self.device + ) + + def init_forward_metadata_capture_cuda_graph( + self, + bs: int, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + forward_mode: ForwardMode, + ): + if forward_mode.is_extend_or_mixed(): + raise NotImplementedError( + f"trtllm_mla CUDA graph capture not supported for {forward_mode}" + ) + + max_blocks = self._calc_padded_blocks(self.max_context_len) + block_kv_indices = self.decode_cuda_graph_kv_indices[:bs, :max_blocks] + + # For capture we don't have req_to_page yet; just zero-fill the block indices. + # The actual indices will be filled on replay. seq_lens_k aliases + # seq_lens_buf (set in init_cuda_graph_state). + metadata = TRTLLMMLADecodeMetadata( + num_extends=0, + block_kv_indices=block_kv_indices, + max_seq_len_k=self.max_context_len, + seq_lens_k=self.cuda_graph_seq_lens_buf[:bs], + ) + + 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, + forward_mode: ForwardMode = None, + req_to_page: torch.Tensor = None, + **kwargs, + ): + if forward_mode is not None and forward_mode.is_extend_or_mixed(): + raise NotImplementedError( + f"trtllm_mla CUDA graph replay not supported for {forward_mode}" + ) + + metadata = self.decode_cuda_graph_metadata[bs] + + # seq_lens_k aliases seq_lens_buf; only block indices need refresh. + # When the buffer is aliased to a peer backend (e.g. drafter aliasing + # the target's kv_indices), the peer's replay has already populated it + # with identical content. + if req_to_page is not None and not self._block_table_aliased: + self._create_block_kv_indices( + bs, + metadata.block_kv_indices.shape[1], + req_pool_indices[:bs], + seq_lens[:bs], + req_to_page, + metadata.block_kv_indices, + ) + + self.forward_decode_metadata = metadata + + def get_cuda_graph_seq_len_fill_value(self): + return 1 + + # ---- Forward: Decode ---- + + def forward_decode( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + layer: PagedAttention, + out_cache_loc: torch.Tensor, + token_to_kv_pool, + bs: int, + save_kv_cache: bool = True, + **kwargs, + ) -> torch.Tensor: + # q is whole Q [T, H, head_dim]; k is whole latent [T, 1, head_dim]. + if save_kv_cache: + assert k is not None + token_to_kv_pool.set_mla_kv_buffer( + layer, + out_cache_loc, + k[..., : self.kv_lora_rank], + k[..., self.kv_lora_rank :], + ) + + metadata = self.forward_decode_metadata + num_extends = metadata.num_extends + q_len_per_req = q.shape[0] // bs if bs > 0 else 1 + + if q_len_per_req > 1 and self.is_draft: + # First draft step catching up its KV after verify: one query entry per token; + # per-token seq_lens advance by 1 so each successive token sees its own KV write. + query = q.view(-1, layer.tp_q_head_num, layer.head_dim).unsqueeze(1) + block_tables = metadata.block_kv_indices[num_extends:].repeat_interleave( + q_len_per_req, dim=0 + ) + base_lens = metadata.seq_lens_k[num_extends:].repeat_interleave( + q_len_per_req + ) + offsets = torch.arange( + q_len_per_req, device=base_lens.device, dtype=base_lens.dtype + ).repeat(bs) + seq_lens = base_lens + offsets + max_seq_len = metadata.max_seq_len_k + q_len_per_req + else: + # Plain decode (q_len=1) or bs-grouped multi-token decode. + query = q.view(bs, -1, layer.tp_q_head_num, layer.head_dim) + block_tables = metadata.block_kv_indices[num_extends:] + seq_lens = metadata.seq_lens_k[num_extends:] + max_seq_len = metadata.max_seq_len_k + + if self.data_type == torch.float8_e4m3fn: + query = query.to(self.data_type) + k_scale = ( + layer.k_scale_float + if getattr(layer, "k_scale_float", None) is not None + else 1.0 + ) + bmm1_scale = k_scale * layer.scaling + else: + bmm1_scale = layer.scaling + + k_cache = token_to_kv_pool.get_key_buffer(layer.layer_id) + if self.data_type != k_cache.dtype: + k_cache = k_cache.to(self.data_type) + kv_cache = k_cache.view(-1, self.page_size, self.kv_cache_dim).unsqueeze(1) + + raw_out = trtllm_batch_decode_with_kv_cache_mla( + query=query, + kv_cache=kv_cache, + workspace_buffer=self.trtllm_workspace, + 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=block_tables, + seq_lens=seq_lens, + max_seq_len=max_seq_len, + bmm1_scale=bmm1_scale, + ) + + return raw_out.view(-1, layer.tp_q_head_num * layer.v_head_dim) + + def forward_extend_chunked( + self, + q, + k, + v, + scaling, + logits_soft_cap, + *, + cum_seq_lens_q, + cum_seq_lens_kv, + max_q_len, + max_kv_len, + seq_lens, + batch_size, + causal, + out: torch.Tensor | None = None, + ): + if causal: + step_counter = getattr(self, "step_counter", None) + if step_counter is not None: + step_counter.record_cache() + + head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + q = q.reshape(-1, self.num_local_heads, head_dim) + k = k.reshape(-1, self.num_local_heads, head_dim) + v = v.reshape(-1, self.num_local_heads, self.v_head_dim) + + # FP8 prefill: if Q is already FP8 (model decided to use FP8 prefill), + # ensure K/V match. If Q is BF16, respect the model's decision. + if q.dtype == torch.float8_e4m3fn: + k = k.to(torch.float8_e4m3fn) + v = v.to(torch.float8_e4m3fn) + + if out is None: + # The ragged path does not support FP8 output. + out_dtype = self.q_data_type + if out_dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + out_dtype = torch.bfloat16 + + out = torch.empty( + q.shape[0], + q.shape[1], + v.shape[2], + device=q.device, + dtype=out_dtype, + ) + + result = trtllm_ragged_attention_deepseek( + query=q, + key=k, + value=v, + workspace_buffer=self.trtllm_workspace, + seq_lens=seq_lens, + max_q_len=max_q_len, + max_kv_len=max_kv_len, + bmm1_scale=scaling, + bmm2_scale=1.0, + o_sf_scale=-1.0, + batch_size=batch_size, + window_left=-1, + cum_seq_lens_q=cum_seq_lens_q, + cum_seq_lens_kv=cum_seq_lens_kv, + enable_pdl=pdl_enabled(), + is_causal=causal, + return_lse=True, + out=out, + ) + + if isinstance(result, tuple): + return result[0], result[1] + return result, None + + +register_backend("trtllm_mla", {AttentionArch.MLA}, TRTLLMMLABackend) diff --git a/python/tokenspeed/runtime/layers/attention/chunk.py b/python/tokenspeed/runtime/layers/attention/chunk.py new file mode 100644 index 0000000..3a4ce7a --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/chunk.py @@ -0,0 +1,209 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from dataclasses import dataclass +from functools import cached_property + +import torch +import triton +import triton.language as tl + +from tokenspeed.runtime.utils import get_colorful_logger +from tokenspeed.runtime.utils.env import global_server_args_dict + +logger = get_colorful_logger(__name__) + + +@triton.jit +def create_chunked_cache_kv_indices_paged( + req_to_page_ptr, # (max_batch, max_pages) + req_pool_indices_ptr, # (batch_size,) + chunk_start_idx_ptr, # (batch_size,) + chunk_seq_lens_ptr, # (batch_size,) + chunk_cum_seq_lens_ptr, # (batch_size + 1,) + chunk_kv_indices_ptr, # (num_chunk_tokens,) + req_to_page_ptr_stride: tl.constexpr, + PAGE_SIZE: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 512 + pid = tl.program_id(axis=0) + + req_pool_index = tl.load(req_pool_indices_ptr + pid) + chunk_kv_indices_offset = tl.load(chunk_cum_seq_lens_ptr + pid) + + chunk_start_pos = tl.load(chunk_start_idx_ptr + pid).to(tl.int32) + chunk_seq_len = tl.load(chunk_seq_lens_ptr + pid).to(tl.int32) + + num_loop = tl.cdiv(chunk_seq_len, BLOCK_SIZE) + for i in range(num_loop): + offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE + mask = offset < chunk_seq_len + token_pos = chunk_start_pos + offset + page_idx = token_pos // PAGE_SIZE + page_id = tl.load( + req_to_page_ptr + req_pool_index * req_to_page_ptr_stride + page_idx, + mask=mask, + ) + kv_slot = page_id * PAGE_SIZE + token_pos % PAGE_SIZE + tl.store( + chunk_kv_indices_ptr + chunk_kv_indices_offset + offset, + kv_slot, + mask=mask, + ) + + +def get_max_chunk_capacity(): + return ( + global_server_args_dict["chunked_prefill_size"] + * global_server_args_dict["mla_chunk_multiplier"] + ) + + +# Here we suppose the length of each chunk is equal +# For example, if we have 4 sequences with seq length [256, 512, 768, 1024], chunk_len = 256 +# num_chunks = cdiv(1024, 256) = 4 +# chunk_starts = [[0, 0, 0, 0], [256, 256, 256, 256], [512, 512, 512, 512], [768, 768, 768, 768]] +# chunk_ends = [[256, 256, 256, 256], [256, 512, 512, 512], [256, 512, 768, 768], [256, 512, 768, 1024]] +# chunk_seq_lens = [[256, 256, 256, 256], [0, 256, 256, 256], [0, 0, 256, 256], [0, 0, 0, 256]] +""" + seq0 seq1 seq2 seq3 +chunk0 -- -- -- -- +chunk1 -- -- -- -- +chunk2 -- -- -- -- +chunk3 -- -- -- -- +""" + + +# starts, ends, len_in_chunk, cum_seq_lens, all satisfy the above layout +@dataclass +class Chunks: + starts: torch.Tensor + ends: torch.Tensor + len_in_chunk: torch.Tensor + + @cached_property + def cum_seq_lens(self): + num_chunks = self.starts.shape[0] + bs = self.starts.shape[1] + result = torch.zeros( + num_chunks, bs + 1, device=self.starts.device, dtype=torch.int32 + ) + torch.cumsum(self.len_in_chunk, dim=1, out=result[:, 1:]) + return result + + +def chunking(prefix_lens: torch.Tensor, num_chunks, batch_size, chunk_len): + starts = ( + torch.arange(num_chunks, device=prefix_lens.device, dtype=torch.int32) + .unsqueeze(1) + .expand(-1, batch_size) + * chunk_len + ) + ends = torch.min(prefix_lens.unsqueeze(0), starts + chunk_len).to(torch.int32) + + chunks = Chunks( + starts=starts, + ends=ends, + len_in_chunk=(ends - starts).clamp(min=0).to(torch.int32), + ) + return chunks + + +def get_chunks_paged( + prefix_lens, prefix_lens_cpu, req_to_page, req_pool_indices, page_size +): + """Page-table aware version of get_chunks.""" + device: torch.device = prefix_lens.device + batch_size = len(prefix_lens_cpu) + + chunk_capacity = get_max_chunk_capacity() + chunk_len = chunk_capacity // batch_size + max_prefix = prefix_lens_cpu.max().item() + num_chunks = (max_prefix + chunk_len - 1) // chunk_len + + chunks = chunking(prefix_lens, num_chunks, batch_size, chunk_len) + chunks_cpu = chunking(prefix_lens_cpu, num_chunks, batch_size, chunk_len) + + num_tokens_per_forward = chunks_cpu.len_in_chunk.sum(dim=1).tolist() + + chunk_kv_indices_list = [] + for idx in range(num_chunks): + chunk_kv_indices = torch.empty( + num_tokens_per_forward[idx], dtype=torch.int32, device=device + ) + create_chunked_cache_kv_indices_paged[(batch_size,)]( + req_to_page, + req_pool_indices, + chunks.starts[idx], + chunks.len_in_chunk[idx], + chunks.cum_seq_lens[idx], + chunk_kv_indices, + req_to_page.shape[1], + page_size, + ) + chunk_kv_indices_list.append(chunk_kv_indices) + + return chunks, chunk_kv_indices_list, chunks_cpu + + +def build_chunked_prefill_metadata_arrays( + extend_prefix_lens, + extend_prefix_lens_cpu, + req_to_page, + req_pool_indices, + page_size, +): + """Build the per-prefix-loop arrays for chunked-prefill MLA. + + Run once per chunked-prefill iteration in the backend's + ``_init_prefill_metadata``. Returns: + + - ``chunked_loop_num``: number of prefix loop iterations + - ``chunk_kv_indices_list``: List[Tensor], paged KV indices per loop_idx + - ``chunked_seq_len``: (chunked_loop_num, num_extends) int32 GPU — per-seq + KV length within each loop_idx (zero for seqs whose prefix doesn't + reach this chunk). + - ``cu_chunked_seq_len``: (chunked_loop_num, num_extends+1) int32 GPU — + cumsum along the seq dim, fed to the chunker as ``cum_seq_lens_kv``. + - ``max_chunk_len_per_loop``: List[int], CPU max-seq-len per loop_idx, + fed to the chunker as ``max_kv_len``. + + The q-side cumsum / max do not appear here: callers alias them to the + causal pass's ``cum_extend_seq_lens`` / ``max_extend_seq_len``, since + every prefix-chunk forward sees the same ``q_lens == extend_seq_lens``. + """ + chunks, chunk_kv_indices_list, chunks_cpu = get_chunks_paged( + extend_prefix_lens, + extend_prefix_lens_cpu, + req_to_page, + req_pool_indices, + page_size, + ) + chunked_loop_num = chunks.starts.shape[0] + max_chunk_len_per_loop = [ + chunks_cpu.len_in_chunk[i].max().item() for i in range(chunked_loop_num) + ] + return ( + chunked_loop_num, + chunk_kv_indices_list, + chunks.len_in_chunk, + chunks.cum_seq_lens, + max_chunk_len_per_loop, + ) diff --git a/python/tokenspeed/runtime/layers/attention/configs/base.py b/python/tokenspeed/runtime/layers/attention/configs/base.py new file mode 100644 index 0000000..3bb2cb3 --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/configs/base.py @@ -0,0 +1,82 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from tokenspeed.runtime.configs.model_config import ModelConfig +from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool +from tokenspeed.runtime.utils.server_args import ServerArgs + + +def resolve_dtype(kv_cache_dtype_str: str) -> torch.dtype: + if kv_cache_dtype_str == "auto": + return torch.bfloat16 + elif kv_cache_dtype_str == "bfloat16": + return torch.bfloat16 + elif kv_cache_dtype_str in ("fp8", "fp8_e4m3"): + return torch.float8_e4m3fn + else: + raise ValueError(f"Unsupported kv_cache_dtype: {kv_cache_dtype_str!r}") + + +@dataclass(kw_only=True) +class BaseAttnConfig: + device: str + backend_name: str + num_attention_heads: int + num_kv_heads: int + head_dim: int + attn_tp_size: int + dtype: torch.dtype + kv_cache_dtype: torch.dtype + page_size: int + context_len: int + max_bs: int + max_graph_bs: int + kv_cache_quant_method: str + speculative_num_steps: int = 0 + speculative_num_draft_tokens: int = 1 + is_draft: bool = False + # DFLASH drafts a whole block in one decode forward (q_len = spec_num_tokens + # per request) instead of Eagle/MTP's per-step single-token decode. Backends + # use this to expand decode metadata to spec_num_tokens rows per request. + draft_block_decode: bool = False + + @classmethod + def generate( + cls, server_args: ServerArgs, model_config: ModelConfig, is_draft: bool = False + ): + raise NotImplementedError("Not Implemented!") + + def cache_cell_size(self) -> int: + raise NotImplementedError("Not Implemented!") + + def create_pool( + self, + num_layers: int, + max_total_num_tokens: int, + rank: int, + enable_memory_saver: bool, + ) -> BaseTokenToKVPool: + raise NotImplementedError("Not Implemented!") diff --git a/python/tokenspeed/runtime/layers/attention/configs/dsa.py b/python/tokenspeed/runtime/layers/attention/configs/dsa.py new file mode 100644 index 0000000..91679df --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/configs/dsa.py @@ -0,0 +1,108 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +from tokenspeed_kernel.platform import current_platform + +from tokenspeed.runtime.configs.model_config import ModelConfig +from tokenspeed.runtime.layers.attention.configs.mla import MLAConfig +from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool +from tokenspeed.runtime.utils.server_args import ServerArgs + +_INDEX_K_FP8_GROUP_SIZE = 128 +_INDEX_K_SCALE_BYTES = torch._utils._element_size(torch.float32) + + +def dsa_index_k_row_bytes(index_head_dim: int) -> int: + if index_head_dim <= 0 or index_head_dim % _INDEX_K_FP8_GROUP_SIZE != 0: + raise ValueError( + f"DSA index_head_dim must be a positive multiple of {_INDEX_K_FP8_GROUP_SIZE}, got {index_head_dim}" + ) + return ( + index_head_dim + + index_head_dim // _INDEX_K_FP8_GROUP_SIZE * _INDEX_K_SCALE_BYTES + ) + + +@dataclass +class DSAConfig(MLAConfig): + index_topk: int + index_head_dim: int + index_n_heads: int + + @classmethod + def generate( + cls, + server_args: ServerArgs, + model_config: ModelConfig, + is_draft: bool = False, + ): + base = MLAConfig.generate(server_args, model_config, is_draft) + if base.kv_cache_dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + platform = current_platform() + if not (platform.is_blackwell_plus or platform.is_cdna4_plus): + raise ValueError( + "GLM DSA FP8 KV cache currently requires NVIDIA Blackwell " + "or AMD CDNA4 sparse attention support; use --kv-cache-dtype " + "auto or bfloat16 on this platform, got " + f"{server_args.kv_cache_dtype}." + ) + return cls( + **base.__dict__, + index_topk=model_config.index_topk, + index_head_dim=model_config.index_head_dim, + index_n_heads=model_config.index_n_heads, + ) + + def cache_cell_size(self) -> int: + index_k_cell_size = dsa_index_k_row_bytes( + self.index_head_dim, + ) + return super().cache_cell_size() + index_k_cell_size + + def create_pool( + self, + num_layers: int, + max_total_num_tokens: int, + rank: int, + enable_memory_saver: bool, + ) -> BaseTokenToKVPool: + from tokenspeed.runtime.layers.attention.kv_cache.dsa import DSATokenToKVPool + + return DSATokenToKVPool( + size=max_total_num_tokens, + dtype=self.kv_cache_dtype, + model_dtype=self.dtype, + quant_method=self.kv_cache_quant_method, + kv_lora_rank=self.kv_lora_rank, + qk_rope_head_dim=self.qk_rope_head_dim, + layer_num=num_layers, + device=self.device, + enable_memory_saver=enable_memory_saver, + max_batch_size=self.max_bs, + max_context_len=self.context_len, + page_size=self.page_size, + rank=rank, + index_head_dim=self.index_head_dim, + ) diff --git a/python/tokenspeed/runtime/layers/attention/configs/mha.py b/python/tokenspeed/runtime/layers/attention/configs/mha.py new file mode 100644 index 0000000..260bb8a --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/configs/mha.py @@ -0,0 +1,169 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from tokenspeed.runtime.configs.model_config import ModelConfig +from tokenspeed.runtime.configs.paged_cache_spec import ( + STATE_LAYER_TYPES, + scheduler_ext_flat_kvcache, +) +from tokenspeed.runtime.layers.attention.configs.base import ( + BaseAttnConfig, + resolve_dtype, +) +from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool +from tokenspeed.runtime.utils.server_args import ServerArgs + + +@dataclass +class MHAConfig(BaseAttnConfig): + # Per-layer attention-type labels + window, forwarded to the KV pool for + # paged_cache_group_specs publication (empty -> single full-history group). + layer_types: tuple[str, ...] = () + sliding_window_tokens: int | tuple[int | None, ...] | None = None + max_scheduled_tokens: int = 0 + # True iff server_args.disaggregation_mode != "null"; the pool's slab + # guards consume it. + pd_disaggregation_enabled: bool = False + # Mamba2/GDN per-state-layer shapes and dtypes (the configs' + # mamba2_cache_params), forwarded to the pool's state slabs. Populated + # only on a flat-built scheduler ext — the radix path keeps its + # SimpleMambaPool state ownership byte-identical (None here means the + # pool neither allocates state slabs nor runs the page-geometry check). + conv_state_shape: tuple[int, ...] | None = None + temporal_state_shape: tuple[int, ...] | None = None + conv_dtype: torch.dtype | None = None + ssm_dtype: torch.dtype | None = None + + @classmethod + def generate( + cls, server_args: ServerArgs, model_config: ModelConfig, is_draft: bool = False + ): + kwargs = {} + if server_args.speculative_algorithm is not None: + kwargs.update( + speculative_num_steps=server_args.speculative_num_steps, + speculative_num_draft_tokens=server_args.speculative_num_draft_tokens, + ) + kv_cache_dtype = server_args.kv_cache_dtype + draft_block_decode = bool( + is_draft and server_args.speculative_algorithm == "DFLASH" + ) + if draft_block_decode: + kv_cache_dtype = "bfloat16" + + hf_config = getattr(model_config, "hf_config", None) + layer_types = tuple(getattr(hf_config, "layer_types", None) or ()) + sliding_window_tokens = getattr(hf_config, "sliding_window", None) + conv_state_shape = temporal_state_shape = None + conv_dtype = ssm_dtype = None + if ( + any(label in STATE_LAYER_TYPES for label in layer_types) + and scheduler_ext_flat_kvcache() + ): + # GDN hybrid on the flat ext: the KV pool owns the recurrent + # state (state slabs), so it needs the mamba2 shapes/dtypes. + # Radix branch untouched: SimpleMambaPool owns the state there. + text_config = getattr(hf_config, "text_config", hf_config) + ( + conv_state_shape, + temporal_state_shape, + conv_dtype, + ssm_dtype, + _, + ) = text_config.mamba2_cache_params + return cls( + device=server_args.device, + context_len=model_config.context_len, + backend_name=( + server_args.attention_backend + if not is_draft + else server_args.drafter_attention_backend + ), + num_attention_heads=model_config.num_attention_heads, + num_kv_heads=model_config.num_key_value_heads, + head_dim=model_config.head_dim, + attn_tp_size=server_args.attn_tp_size or server_args.mapping.attn.tp_size, + dtype=model_config.dtype, + kv_cache_dtype=resolve_dtype(kv_cache_dtype), + page_size=server_args.block_size, + max_bs=server_args.max_num_seqs + // (server_args.data_parallel_size or server_args.mapping.attn.dp_size), + max_graph_bs=server_args.max_cudagraph_capture_size, + kv_cache_quant_method=server_args.kv_cache_quant_method, + is_draft=is_draft, + draft_block_decode=draft_block_decode, + layer_types=layer_types, + sliding_window_tokens=sliding_window_tokens, + max_scheduled_tokens=getattr(server_args, "chunked_prefill_size", 8192), + pd_disaggregation_enabled=getattr( + server_args, "disaggregation_mode", "null" + ) + != "null", + conv_state_shape=conv_state_shape, + temporal_state_shape=temporal_state_shape, + conv_dtype=conv_dtype, + ssm_dtype=ssm_dtype, + **kwargs, + ) + + def cache_cell_size(self) -> int: + return ( + max(self.num_kv_heads // self.attn_tp_size, 1) + * self.head_dim + * 2 + * torch._utils._element_size(self.kv_cache_dtype) + ) + + def create_pool( + self, + num_layers: int, + max_total_num_tokens: int, + rank: int, + enable_memory_saver: bool, + ) -> BaseTokenToKVPool: + from tokenspeed.runtime.layers.attention.kv_cache.mha import MHATokenToKVPool + + return MHATokenToKVPool( + size=max_total_num_tokens, + dtype=self.kv_cache_dtype, + head_num=max(self.num_kv_heads // self.attn_tp_size, 1), + head_dim=self.head_dim, + layer_num=num_layers, + device=self.device, + enable_memory_saver=enable_memory_saver, + max_batch_size=self.max_bs, + max_context_len=self.context_len, + page_size=self.page_size, + rank=rank, + layer_types=self.layer_types, + sliding_window_tokens=self.sliding_window_tokens, + max_scheduled_tokens=self.max_scheduled_tokens, + pd_disaggregation_enabled=self.pd_disaggregation_enabled, + conv_state_shape=self.conv_state_shape, + temporal_state_shape=self.temporal_state_shape, + conv_dtype=self.conv_dtype, + ssm_dtype=self.ssm_dtype, + ) diff --git a/python/tokenspeed/runtime/layers/attention/configs/mla.py b/python/tokenspeed/runtime/layers/attention/configs/mla.py new file mode 100644 index 0000000..28ae592 --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/configs/mla.py @@ -0,0 +1,120 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from tokenspeed.runtime.configs.model_config import ModelConfig +from tokenspeed.runtime.layers.attention.configs.base import ( + BaseAttnConfig, + resolve_dtype, +) +from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool +from tokenspeed.runtime.utils.server_args import ServerArgs + + +@dataclass +class MLAConfig(BaseAttnConfig): + kv_lora_rank: int + qk_nope_head_dim: int + qk_rope_head_dim: int + v_head_dim: int + scaling: float + kv_cache_dim: int + + @classmethod + def generate( + cls, server_args: ServerArgs, model_config: ModelConfig, is_draft: bool = False + ): + kwargs = {} + if server_args.speculative_algorithm is not None: + kwargs.update( + speculative_num_steps=server_args.speculative_num_steps, + speculative_num_draft_tokens=server_args.speculative_num_draft_tokens, + ) + return cls( + device=server_args.device, + context_len=model_config.context_len, + backend_name=( + server_args.attention_backend + if not is_draft + else server_args.drafter_attention_backend + ), + num_attention_heads=model_config.num_attention_heads, + num_kv_heads=model_config.num_key_value_heads, + head_dim=model_config.head_dim, + attn_tp_size=server_args.attn_tp_size or server_args.mapping.attn.tp_size, + dtype=model_config.dtype, + kv_cache_dtype=resolve_dtype(server_args.kv_cache_dtype), + page_size=server_args.block_size, + max_graph_bs=server_args.max_cudagraph_capture_size, + max_bs=server_args.max_num_seqs + // (server_args.data_parallel_size or server_args.mapping.attn.dp_size), + kv_cache_quant_method=server_args.kv_cache_quant_method, + is_draft=is_draft, + kv_lora_rank=model_config.kv_lora_rank, + qk_nope_head_dim=model_config.qk_nope_head_dim, + qk_rope_head_dim=model_config.qk_rope_head_dim, + v_head_dim=model_config.v_head_dim, + scaling=model_config.scaling, + kv_cache_dim=model_config.kv_lora_rank + model_config.qk_rope_head_dim, + **kwargs, + ) + + def cache_cell_size(self) -> int: + if self.kv_cache_quant_method == "per_token_head": + cell_size = ( + self.kv_lora_rank * torch._utils._element_size(self.kv_cache_dtype) + + self.qk_rope_head_dim * torch._utils._element_size(self.dtype) + + 1 * torch._utils._element_size(torch.float32) + ) + else: + cell_size = ( + self.kv_lora_rank + self.qk_rope_head_dim + ) * torch._utils._element_size(self.kv_cache_dtype) + return cell_size + + def create_pool( + self, + num_layers: int, + max_total_num_tokens: int, + rank: int, + enable_memory_saver: bool, + ) -> BaseTokenToKVPool: + from tokenspeed.runtime.layers.attention.kv_cache.mla import MLATokenToKVPool + + return MLATokenToKVPool( + size=max_total_num_tokens, + dtype=self.kv_cache_dtype, + model_dtype=self.dtype, + quant_method=self.kv_cache_quant_method, + kv_lora_rank=self.kv_lora_rank, + qk_rope_head_dim=self.qk_rope_head_dim, + layer_num=num_layers, + device=self.device, + enable_memory_saver=enable_memory_saver, + max_batch_size=self.max_bs, + max_context_len=self.context_len, + page_size=self.page_size, + rank=rank, + ) diff --git a/python/tokenspeed/runtime/layers/attention/deepseek_v4/metadata.py b/python/tokenspeed/runtime/layers/attention/deepseek_v4/metadata.py new file mode 100644 index 0000000..38509e5 --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/deepseek_v4/metadata.py @@ -0,0 +1,159 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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 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. + +from __future__ import annotations + +from dataclasses import dataclass, field + +import torch + +from tokenspeed.runtime.execution.forward_batch_info import ForwardMode +from tokenspeed.runtime.layers.attention.kv_cache.deepseek_v4 import ( + DeepseekV4CacheMetadata, +) + + +@dataclass +class DeepseekV4IndexerPrefillChunkPlan: + token_start: int + token_end: int + request_start: int + request_end: int + slot_start: int + slot_end: int + gather_row_start: int + gather_row_end: int + max_seq_len_k: int + cu_seq_lens_start: int + cu_seq_lens_end: int + skip_kv_gather: bool = False + + +@dataclass +class DeepseekV4IndexerPrefillMetadata: + chunks: tuple[DeepseekV4IndexerPrefillChunkPlan, ...] + chunk_specs: torch.Tensor + chunk_offsets: torch.Tensor + slots: torch.Tensor + cu_seq_lens: torch.Tensor + cu_seqlen_k_start: torch.Tensor + cu_seqlen_k_end: torch.Tensor + seq_lens_k: torch.Tensor + + def max_gather_rows(self) -> int: + if not self.chunks: + return 0 + return max(max(0, chunk.slot_end - chunk.slot_start) for chunk in self.chunks) + + @classmethod + def empty(cls, device: torch.device) -> "DeepseekV4IndexerPrefillMetadata": + return cls( + chunks=(), + chunk_specs=torch.empty((0, 5), dtype=torch.int64, device="cpu"), + chunk_offsets=torch.empty((0, 7), dtype=torch.int64, device="cpu"), + slots=torch.empty(0, dtype=torch.int64, device=device), + cu_seq_lens=torch.empty(0, dtype=torch.int32, device=device), + cu_seqlen_k_start=torch.empty(0, dtype=torch.int32, device=device), + cu_seqlen_k_end=torch.empty(0, dtype=torch.int32, device=device), + seq_lens_k=torch.empty(0, dtype=torch.int32, device=device), + ) + + +@dataclass +class DeepseekV4IndexerDecodePlan: + context_lens: torch.Tensor + block_table: torch.Tensor + max_context_len: int + + +@dataclass +class DeepseekV4IndexerBatchMetadata: + positions: torch.Tensor + token_to_req_indices: torch.Tensor + seq_lens_cpu: torch.Tensor + query_lens_cpu: torch.Tensor + num_prefill_tokens: int + num_decode_tokens: int + + +@dataclass +class DeepseekV4AttentionMetadata: + decode_swa_indices: torch.Tensor | None = None + decode_swa_lens: torch.Tensor | None = None + decode_swa_window_size: int = 0 + decode_swa_block_size: int = 0 + # Cache for dense compressed decode attention indices/lens. CSA decode uses + # dynamic top-k indices and does not populate this cache. + decode_dense_compressed_indices_cache: dict[ + tuple[int, int, int, int], tuple[torch.Tensor, torch.Tensor] + ] = field(default_factory=dict) + decode_dense_compressed_indices_capture_safe_keys: set[ + tuple[int, int, int, int] + ] = field(default_factory=set) + + +@dataclass +class DeepseekV4IndexerMetadata: + decode_schedule_metadata_cache: dict[tuple[int, int, int], torch.Tensor] = field( + default_factory=dict + ) + decode_plan_cache: dict[tuple[int, int, int], DeepseekV4IndexerDecodePlan] = field( + default_factory=dict + ) + decode_plan_refreshed_keys: set[tuple[int, int, int]] = field(default_factory=set) + prefill_plan_cache: dict[tuple[int, int, int], DeepseekV4IndexerPrefillMetadata] = ( + field(default_factory=dict) + ) + + +@dataclass +class DeepseekV4SparseIndexerMetadata: + batch_metadata: DeepseekV4IndexerBatchMetadata | None = None + prefill_metadata: DeepseekV4IndexerPrefillMetadata | None = None + decode_plan: DeepseekV4IndexerDecodePlan | None = None + decode_schedule_metadata: torch.Tensor | None = None + + +@dataclass +class DeepseekV4ForwardMetadata: + req_pool_indices: torch.Tensor + seq_lens: torch.Tensor + query_lens: torch.Tensor + query_start_loc: torch.Tensor + token_to_req_indices: torch.Tensor + cache: DeepseekV4CacheMetadata + attention: DeepseekV4AttentionMetadata = field( + default_factory=DeepseekV4AttentionMetadata + ) + indexer: DeepseekV4IndexerMetadata = field( + default_factory=DeepseekV4IndexerMetadata + ) + forward_mode: ForwardMode | None = None + # Padding mask for CUDA graph replay rows; this is not mixed-batch state. + is_valid_token: torch.Tensor | None = None + # CPU lens are retained for sparse prefill/indexer planning without + # forcing another device-to-host sync in the model path. + seq_lens_cpu: torch.Tensor | None = None + query_lens_cpu: torch.Tensor | None = None + # Cached split boundary derived from scheduler num_extends/query_lens. + num_prefill_reqs: int = 0 + num_prefill_tokens: int = 0 + + def decode_req_count(self) -> int: + return max(0, int(self.req_pool_indices.shape[0]) - int(self.num_prefill_reqs)) + + def decode_token_count(self) -> int: + return max( + 0, + int(self.token_to_req_indices.shape[0]) - int(self.num_prefill_tokens), + ) diff --git a/python/tokenspeed/runtime/layers/attention/deepseek_v4_ops.py b/python/tokenspeed/runtime/layers/attention/deepseek_v4_ops.py new file mode 100644 index 0000000..f86ed47 --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/deepseek_v4_ops.py @@ -0,0 +1,1085 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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 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. +# +# DeepSeek V4 attention helpers keep runtime validation here; production Triton +# kernels live under tokenspeed-kernel ops. + +"""DeepSeek V4 attention kernel boundaries. + +Keep the model layer independent from the CUDA extension import details. The +runtime requires TokenSpeed's own built DeepSeek V4 attention op. +""" + +from __future__ import annotations + +import math + +import torch +from tokenspeed_kernel.ops.attention.cuda.deepseek_v4 import ( + fused_qnorm_rope_kv_insert as _cuda_fused_qnorm_rope_kv_insert, +) +from tokenspeed_kernel.ops.attention.cuda.deepseek_v4 import ( + has_fused_qnorm_rope_kv_insert as _cuda_has_fused_qnorm_rope_kv_insert, +) +from tokenspeed_kernel.ops.attention.triton.deepseek_v4 import ( + deepseek_v4_build_dense_prefill_local_compressed_indices, + deepseek_v4_combine_dense_swa_indices, + deepseek_v4_combine_topk_swa_indices, + deepseek_v4_compressed_slot_mapping, + deepseek_v4_compute_global_topk_indices_and_lens, + deepseek_v4_decode_swa_indices_and_lens, + deepseek_v4_dequantize_and_gather_k_cache, +) +from tokenspeed_kernel.ops.attention.triton.deepseek_v4 import ( + deepseek_v4_fused_csa_indexer_mxfp4_cache_insert as _triton_fused_csa_indexer_mxfp4_cache_insert, +) +from tokenspeed_kernel.ops.attention.triton.deepseek_v4 import ( + deepseek_v4_fused_indexer_q_rope_hadamard_mxfp4 as _triton_fused_indexer_q_rope_hadamard_mxfp4, +) +from tokenspeed_kernel.ops.attention.triton.deepseek_v4 import ( + deepseek_v4_fused_inv_rope_fp8_quant, +) +from tokenspeed_kernel.ops.attention.triton.deepseek_v4 import ( + deepseek_v4_fused_sparse_compress_cache_insert as _triton_fused_sparse_compress_cache_insert, +) +from tokenspeed_kernel.ops.attention.triton.deepseek_v4 import ( + deepseek_v4_save_compressor_state as _triton_save_compressor_state, +) +from tokenspeed_kernel.ops.attention.triton.deepseek_v4 import ( + write_deepseek_v4_indexer_mxfp4_cache_cuda as _triton_write_indexer_mxfp4_cache_cuda, +) +from tokenspeed_kernel.ops.transform import hadamard_transform + +from tokenspeed.runtime.configs.deepseek_v4_cache_spec import ( + DEEPSEEK_V4_FP8_MAX, + DEEPSEEK_V4_FP8_QUANT_BLOCK, + DEEPSEEK_V4_MXFP4_BLOCK_SIZE, + deepseek_v4_indexer_fp8_layout_from_row_bytes, + deepseek_v4_indexer_fp8_row_bytes, + deepseek_v4_indexer_fp8_scale_bytes, + deepseek_v4_indexer_mxfp4_layout_from_row_bytes, + deepseek_v4_indexer_mxfp4_row_bytes, + deepseek_v4_nope_dim, + deepseek_v4_swa_row_bytes, + deepseek_v4_swa_scale_dim, + deepseek_v4_swa_token_stride, +) + +__all__ = ( + "deepseek_v4_build_dense_prefill_local_compressed_indices", + "deepseek_v4_combine_dense_swa_indices", + "deepseek_v4_combine_topk_swa_indices", + "deepseek_v4_compressed_slot_mapping", + "deepseek_v4_compute_global_topk_indices_and_lens", + "deepseek_v4_decode_swa_indices_and_lens", + "deepseek_v4_dequantize_and_gather_k_cache", + "deepseek_v4_fused_inv_rope_fp8_quant", + "deepseek_v4_csa_compress_kv_cache_insert", + "deepseek_v4_csa_indexer_cache_insert", + "deepseek_v4_hca_compress_kv_cache_insert", + "deepseek_v4_prepare_indexer_q_mxfp4", + "dequantize_deepseek_v4_fp8_ds_mla_cache", + "fused_qnorm_rope_kv_insert", + "read_deepseek_v4_indexer_fp8_cache", + "read_deepseek_v4_indexer_mxfp4_cache", + "save_deepseek_v4_compressor_state", + "write_deepseek_v4_indexer_fp8_cache", + "write_deepseek_v4_indexer_mxfp4_cache", +) + + +def _indexer_mxfp4_layout_from_cache( + cache_2d: torch.Tensor, + block_size: int, +) -> tuple[int, int, int]: + if cache_2d.dim() != 2: + raise ValueError(f"cache_2d must be 2-D, got {tuple(cache_2d.shape)}") + row_bytes = cache_2d.shape[1] // block_size + if cache_2d.shape[1] % block_size != 0: + raise ValueError( + "MXFP4 indexer cache row size must match value+scale layout, " + f"got cache shape {tuple(cache_2d.shape)} and block_size={block_size}" + ) + return deepseek_v4_indexer_mxfp4_layout_from_row_bytes(row_bytes) + + +def _indexer_fp8_layout_from_cache( + cache_2d: torch.Tensor, + block_size: int, +) -> tuple[int, int]: + if cache_2d.dim() != 2: + raise ValueError(f"cache_2d must be 2-D, got {tuple(cache_2d.shape)}") + row_bytes = cache_2d.shape[1] // block_size + if cache_2d.shape[1] % block_size != 0: + raise ValueError( + "FP8 indexer cache row size must match value+scale layout, " + f"got cache shape {tuple(cache_2d.shape)} and block_size={block_size}" + ) + return deepseek_v4_indexer_fp8_layout_from_row_bytes(row_bytes) + + +def fused_qnorm_rope_kv_insert( + q: torch.Tensor, + kv: torch.Tensor, + swa_kv_cache_2d: torch.Tensor, + slot_mapping: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + rms_norm_eps: float, + block_size: int, +) -> None: + """Run the DeepSeek V4 fused SWA cache insert op. + + Expected contract: + - q: [tokens, local_heads, 512], mutated in place by RMSNorm/RoPE + - kv: [tokens, 512], source KV latent before RoPE/quant insert + - swa_kv_cache_2d: uint8 cache blocks flattened as [num_blocks, block_bytes] + - slot_mapping: output token slots in the paged SWA cache + - positions: absolute token positions + """ + + if not _cuda_has_fused_qnorm_rope_kv_insert(): + raise RuntimeError( + "DeepSeek V4 fused SWA cache insert op " + "fused_deepseek_v4_qnorm_rope_kv_rope_quant_insert is unavailable. " + "Build `tokenspeed-kernel/python` so the deepseek_v4_attention CUDA " + "library is present before running this path." + ) + + _cuda_fused_qnorm_rope_kv_insert( + q, + kv, + swa_kv_cache_2d, + slot_mapping, + positions.to(torch.int64), + cos_sin_cache, + rms_norm_eps, + block_size, + ) + + +def _apply_gptj_rope_tail_rows( + x: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + rope_dim: int, +) -> torch.Tensor: + out = x.float().clone() + half_rope = rope_dim // 2 + nope_dim = x.shape[-1] - rope_dim + cos = cos_sin_cache[positions.long(), :half_rope].float() + sin = cos_sin_cache[positions.long(), half_rope:rope_dim].float() + even = out[..., nope_dim::2].clone() + odd = out[..., nope_dim + 1 :: 2].clone() + while cos.ndim < even.ndim: + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + out[..., nope_dim::2] = even * cos - odd * sin + out[..., nope_dim + 1 :: 2] = even * sin + odd * cos + return out + + +def _fp8_e4m3_pow2_bytes(block: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + scale = max(float(block.detach().abs().max()) / DEEPSEEK_V4_FP8_MAX, 1.0e-10) + scale = 2.0 ** math.ceil(math.log2(scale)) + scaled = torch.clamp(block / scale, -DEEPSEEK_V4_FP8_MAX, DEEPSEEK_V4_FP8_MAX) + return scaled.to(torch.float8_e4m3fn).view(torch.uint8), block.new_tensor(scale) + + +def _e2m1_values(nibbles: torch.Tensor) -> torch.Tensor: + table = nibbles.new_tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32 + ) + magnitude = table[(nibbles & 0x7).long()] + sign = torch.where((nibbles & 0x8) != 0, -1.0, 1.0) + return magnitude * sign + + +def _deepseek_v4_hadamard_rotate(x: torch.Tensor) -> torch.Tensor: + shape = x.shape + rotated = hadamard_transform( + x.to(torch.bfloat16).reshape(-1, shape[-1]).contiguous(), + scale=shape[-1] ** -0.5, + ) + return rotated.reshape(shape) + + +def dequantize_deepseek_v4_fp8_ds_mla_cache( + cache_2d: torch.Tensor, + slot_mapping: torch.Tensor, + block_size: int = 64, + *, + head_dim: int, + rope_dim: int, +) -> torch.Tensor: + """Dequantize DeepSeek V4 `fp8_ds_mla` rows selected by global slots.""" + + nope_dim = deepseek_v4_nope_dim(head_dim, rope_dim) + token_stride = deepseek_v4_swa_token_stride(head_dim, rope_dim) + scale_dim = deepseek_v4_swa_scale_dim(head_dim, rope_dim) + min_stride = block_size * (token_stride + scale_dim) + if cache_2d.dtype != torch.uint8: + raise TypeError(f"cache_2d must be uint8, got {cache_2d.dtype}") + if cache_2d.dim() != 2 or cache_2d.shape[1] < min_stride: + raise ValueError( + f"cache_2d must be [pages, >= {min_stride}], got {tuple(cache_2d.shape)}" + ) + + out_shape = (slot_mapping.numel(), head_dim) + if slot_mapping.numel() == 0: + return torch.empty(out_shape, device=cache_2d.device, dtype=torch.bfloat16) + + flat_cache = cache_2d.reshape(-1) + num_nope_blocks = nope_dim // DEEPSEEK_V4_FP8_QUANT_BLOCK + + slots = slot_mapping.to(torch.int64) + valid = slots >= 0 + safe_slots = torch.where(valid, slots, torch.zeros_like(slots)) + pages = torch.div(safe_slots, block_size, rounding_mode="floor") + pos = safe_slots % block_size + page_base = pages * cache_2d.stride(0) + value_base = page_base + pos * token_stride + scale_base = page_base + block_size * token_stride + pos * scale_dim + + value_offsets = ( + value_base[:, None] + + torch.arange(token_stride, device=cache_2d.device, dtype=torch.int64)[None, :] + ) + row_bytes = flat_cache[value_offsets] + nope = row_bytes[:, :nope_dim].contiguous().view(torch.float8_e4m3fn) + + scale_offsets = ( + scale_base[:, None] + + torch.arange(num_nope_blocks, device=cache_2d.device, dtype=torch.int64)[ + None, : + ] + ) + scales = torch.pow(2.0, flat_cache[scale_offsets].to(torch.int32) - 127) + scales = scales.float().repeat_interleave(DEEPSEEK_V4_FP8_QUANT_BLOCK, dim=1) + + rope = row_bytes[:, nope_dim:token_stride].contiguous() + out = torch.cat([nope.float() * scales, rope.view(torch.bfloat16).float()], dim=1) + out = out.to(torch.bfloat16) + return torch.where(valid[:, None], out, torch.zeros_like(out)) + + +def deepseek_v4_prepare_indexer_q_mxfp4( + index_q: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + weights: torch.Tensor, + softmax_scale: float, + head_scale: float, +) -> tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor]: + """Apply indexer Q RoPE and return DeepGEMM-ready MXFP4 values/scales.""" + + if index_q.dim() != 3: + raise ValueError(f"index_q must be [tokens, heads, dim], got {index_q.shape}") + if index_q.shape[-1] % DEEPSEEK_V4_MXFP4_BLOCK_SIZE != 0: + raise ValueError( + "MXFP4 index_q dim must be divisible by " + f"{DEEPSEEK_V4_MXFP4_BLOCK_SIZE}, got {index_q.shape[-1]}" + ) + rope_dim = int(cos_sin_cache.shape[-1]) + if index_q.shape[-1] <= rope_dim: + raise ValueError( + f"index_q dim must be larger than rope_dim={rope_dim}, got {index_q.shape}" + ) + if weights.dim() == 3: + weights = weights.squeeze(-1) + if weights.shape != index_q.shape[:2]: + raise ValueError(f"weights must be [tokens, heads], got {tuple(weights.shape)}") + if not index_q.is_cuda: + raise ValueError( + "deepseek_v4_prepare_indexer_q_mxfp4 only supports CUDA tensors." + ) + return _triton_fused_indexer_q_rope_hadamard_mxfp4( + index_q=index_q, + positions=positions, + cos_sin_cache=cos_sin_cache, + weights=weights, + softmax_scale=softmax_scale, + head_scale=head_scale, + ) + + +def _fp8_ds_mla_cache_rows( + normed: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + compress_ratio: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + num_rows = normed.shape[0] + head_dim = int(normed.shape[-1]) + rope_dim = int(cos_sin_cache.shape[-1]) + nope_dim = deepseek_v4_nope_dim(head_dim, rope_dim) + scale_dim = deepseek_v4_swa_scale_dim(head_dim, rope_dim) + quant_input = normed.to(torch.bfloat16).float() + nope_blocks = quant_input[:, :nope_dim].reshape( + num_rows, + nope_dim // DEEPSEEK_V4_FP8_QUANT_BLOCK, + DEEPSEEK_V4_FP8_QUANT_BLOCK, + ) + absmax = nope_blocks.detach().abs().amax(dim=-1).clamp_min(1.0e-4) + exponent = torch.ceil(torch.log2(absmax / DEEPSEEK_V4_FP8_MAX)) + scaled = torch.clamp( + nope_blocks * torch.pow(2.0, -exponent).unsqueeze(-1), + -DEEPSEEK_V4_FP8_MAX, + DEEPSEEK_V4_FP8_MAX, + ) + value_bytes = ( + scaled.to(torch.float8_e4m3fn) + .view(torch.uint8) + .reshape( + num_rows, + nope_dim, + ) + ) + scale_bytes = torch.clamp(exponent + 127.0, 0.0, 255.0).to(torch.uint8) + scale_pad = scale_dim - scale_bytes.shape[1] + if scale_pad > 0: + scale_bytes = torch.cat( + [scale_bytes, torch.zeros_like(scale_bytes[:, :scale_pad])], + dim=-1, + ) + + compressed_positions = ( + torch.div(positions.to(torch.int64), compress_ratio, rounding_mode="floor") + * compress_ratio + ) + rotated = _apply_gptj_rope_tail_rows( + normed, + compressed_positions, + cos_sin_cache, + rope_dim, + ).to(torch.bfloat16) + rope_bytes = rotated[:, nope_dim:].contiguous().view(torch.uint8) + rope_bytes = rope_bytes.reshape(num_rows, rope_dim * 2) + return value_bytes, scale_bytes, rope_bytes + + +def _write_fp8_ds_mla_cache_rows_capturable( + normed: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + kv_cache_2d: torch.Tensor, + kv_slot_mapping: torch.Tensor, + valid: torch.Tensor, + kv_cache_block_size: int, + compress_ratio: int, +) -> None: + num_rows = normed.shape[0] + if num_rows == 0: + return + + head_dim = int(normed.shape[-1]) + rope_dim = int(cos_sin_cache.shape[-1]) + nope_dim = deepseek_v4_nope_dim(head_dim, rope_dim) + token_stride = deepseek_v4_swa_token_stride(head_dim, rope_dim) + scale_dim = deepseek_v4_swa_scale_dim(head_dim, rope_dim) + slots = kv_slot_mapping[:num_rows].to(torch.int64) + valid = valid[:num_rows] & (slots >= 0) + if not (slots.is_cuda and torch.cuda.is_current_stream_capturing()): + if not bool(valid.any()): + return + normed = normed[:num_rows][valid] + positions = positions[:num_rows][valid] + slots = slots[valid] + valid = torch.ones_like(slots, dtype=torch.bool) + num_rows = slots.numel() + safe_slots = torch.where(valid, slots, torch.zeros_like(slots)) + block_idx = torch.div(safe_slots, kv_cache_block_size, rounding_mode="floor") + pos_in_block = safe_slots % kv_cache_block_size + block_base = block_idx * kv_cache_2d.stride(0) + token_base = block_base + pos_in_block * token_stride + scale_base = ( + block_base + kv_cache_block_size * token_stride + pos_in_block * scale_dim + ) + + value_bytes, scale_bytes, rope_bytes = _fp8_ds_mla_cache_rows( + normed[:num_rows], positions[:num_rows], cos_sin_cache, compress_ratio + ) + flat_cache = kv_cache_2d.reshape(-1) + value_offsets = ( + token_base[:, None] + + torch.arange( + nope_dim, + device=kv_cache_2d.device, + dtype=torch.int64, + )[None, :] + ) + scale_offsets = ( + scale_base[:, None] + + torch.arange( + scale_dim, + device=kv_cache_2d.device, + dtype=torch.int64, + )[None, :] + ) + rope_offsets = ( + token_base[:, None] + + nope_dim + + torch.arange( + rope_dim * 2, + device=kv_cache_2d.device, + dtype=torch.int64, + )[None, :] + ) + flat_cache[value_offsets] = value_bytes + flat_cache[scale_offsets] = scale_bytes + flat_cache[rope_offsets] = rope_bytes + + +def save_deepseek_v4_compressor_state( + kv: torch.Tensor, + score: torch.Tensor, + ape: torch.Tensor, + state_cache: torch.Tensor, + slot_mapping: torch.Tensor, + positions: torch.Tensor, + block_size: int, + compress_ratio: int, +) -> None: + """Save DeepSeek V4 compressor residual state into paged SWA-style cache. + + This correctness-first state write packs `[kv_state, score_state]`, each + with width `coff * head_dim`; score state includes the APE row selected by + `position % compress_ratio`. + """ + + if kv.shape != score.shape: + raise ValueError( + f"kv and score shapes must match, got {kv.shape} vs {score.shape}" + ) + if kv.dim() != 2: + raise ValueError(f"kv/score must be [tokens, state_width], got {kv.shape}") + if state_cache.dim() != 3: + raise ValueError( + "state_cache must be [blocks, block_size, 2 * state_width], " + f"got {state_cache.shape}" + ) + if block_size != state_cache.shape[1]: + raise ValueError( + f"block_size={block_size} does not match " + f"state_cache.shape[1]={state_cache.shape[1]}" + ) + state_width = kv.shape[-1] + if state_cache.shape[-1] != state_width * 2: + raise ValueError( + f"state_cache last dim must be {state_width * 2}, " + f"got {state_cache.shape[-1]}" + ) + if ape.shape != (compress_ratio, state_width): + raise ValueError( + f"ape must be [{compress_ratio}, {state_width}], got {tuple(ape.shape)}" + ) + + num_actual = min(slot_mapping.numel(), kv.shape[0]) + if num_actual == 0: + return + if not state_cache.is_cuda: + raise ValueError( + "save_deepseek_v4_compressor_state only supports CUDA tensors." + ) + + _triton_save_compressor_state( + kv=kv, + score=score, + ape=ape, + state_cache=state_cache, + slot_mapping=slot_mapping, + positions=positions, + block_size=block_size, + compress_ratio=compress_ratio, + ) + + +def write_deepseek_v4_indexer_fp8_cache( + index_k: torch.Tensor, + cache_2d: torch.Tensor, + slot_mapping: torch.Tensor, + block_size: int = 64, +) -> None: + """Write FP8 indexer keys using `[values | fp32 scale]` page layout.""" + + if index_k.dim() != 2: + raise ValueError(f"index_k must be [tokens, dim], got {tuple(index_k.shape)}") + index_head_dim = int(index_k.shape[-1]) + scale_bytes = deepseek_v4_indexer_fp8_scale_bytes(index_head_dim) + row_bytes = deepseek_v4_indexer_fp8_row_bytes(index_head_dim) + if cache_2d.dtype != torch.uint8: + raise TypeError(f"cache_2d must be uint8, got {cache_2d.dtype}") + min_stride = block_size * row_bytes + if cache_2d.dim() != 2 or cache_2d.shape[1] < min_stride: + raise ValueError( + f"cache_2d must be [pages, >= {min_stride}], " + f"got {tuple(cache_2d.shape)}" + ) + + flat_cache = cache_2d.reshape(-1) + num_actual = min(slot_mapping.numel(), index_k.shape[0]) + for token_idx in range(num_actual): + slot = int(slot_mapping[token_idx].item()) + if slot < 0: + continue + page = slot // block_size + pos = slot % block_size + page_base = page * cache_2d.stride(0) + value_base = page_base + pos * index_head_dim + scale_base = page_base + block_size * index_head_dim + pos * scale_bytes + q_bytes, scale = _fp8_e4m3_pow2_bytes(index_k[token_idx].float()) + flat_cache[value_base : value_base + index_head_dim].copy_(q_bytes) + flat_cache[scale_base : scale_base + scale_bytes].copy_( + scale.reshape(1).view(torch.uint8) + ) + + +def write_deepseek_v4_indexer_mxfp4_cache( + index_k: torch.Tensor, + cache_2d: torch.Tensor, + slot_mapping: torch.Tensor, + block_size: int = 64, +) -> None: + """Write MXFP4 indexer keys using the `[values | ue8m0 scales]` layout.""" + + if index_k.dim() != 2: + raise ValueError(f"index_k must be [tokens, dim], got {tuple(index_k.shape)}") + index_head_dim = int(index_k.shape[-1]) + row_bytes = deepseek_v4_indexer_mxfp4_row_bytes(index_head_dim) + if cache_2d.dtype != torch.uint8: + raise TypeError(f"cache_2d must be uint8, got {cache_2d.dtype}") + min_stride = block_size * row_bytes + if cache_2d.dim() != 2 or cache_2d.shape[1] < min_stride: + raise ValueError( + f"cache_2d must be [pages, >= {min_stride}], got {tuple(cache_2d.shape)}" + ) + + num_actual = min(slot_mapping.numel(), index_k.shape[0]) + if num_actual == 0: + return + if not index_k.is_cuda: + raise ValueError( + "write_deepseek_v4_indexer_mxfp4_cache only supports CUDA tensors." + ) + valid = torch.ones(num_actual, device=index_k.device, dtype=torch.bool) + _triton_write_indexer_mxfp4_cache_cuda( + index_k[:num_actual], + cache_2d, + slot_mapping[:num_actual], + valid, + block_size, + ) + + +def _write_deepseek_v4_indexer_fp8_cache_capturable( + index_k: torch.Tensor, + cache_2d: torch.Tensor, + slot_mapping: torch.Tensor, + valid: torch.Tensor, + block_size: int = 64, +) -> None: + num_rows = min(slot_mapping.numel(), index_k.shape[0]) + if num_rows == 0: + return + + index_head_dim = int(index_k.shape[-1]) + scale_bytes = deepseek_v4_indexer_fp8_scale_bytes(index_head_dim) + rows = index_k[:num_rows].float() + scale = (rows.detach().abs().amax(dim=-1) / DEEPSEEK_V4_FP8_MAX).clamp_min(1.0e-10) + scale = torch.pow(2.0, torch.ceil(torch.log2(scale))) + value_bytes = ( + torch.clamp( + rows / scale.unsqueeze(-1), + -DEEPSEEK_V4_FP8_MAX, + DEEPSEEK_V4_FP8_MAX, + ) + .to(torch.float8_e4m3fn) + .view(torch.uint8) + ) + + slots = slot_mapping[:num_rows].to(torch.int64) + valid = valid[:num_rows] & (slots >= 0) + if not (slots.is_cuda and torch.cuda.is_current_stream_capturing()): + if not bool(valid.any()): + return + rows = rows[valid] + slots = slots[valid] + scale = scale[valid] + value_bytes = value_bytes[valid] + valid = torch.ones_like(slots, dtype=torch.bool) + num_rows = slots.numel() + safe_slots = torch.where(valid, slots, torch.zeros_like(slots)) + pages = torch.div(safe_slots, block_size, rounding_mode="floor") + pos = safe_slots % block_size + page_base = pages * cache_2d.stride(0) + value_base = page_base + pos * index_head_dim + scale_base = page_base + block_size * index_head_dim + pos * scale_bytes + + flat_cache = cache_2d.reshape(-1) + value_offsets = ( + value_base[:, None] + + torch.arange( + index_head_dim, + device=cache_2d.device, + dtype=torch.int64, + )[None, :] + ) + scale_offsets = ( + scale_base[:, None] + + torch.arange(scale_bytes, device=cache_2d.device, dtype=torch.int64)[None, :] + ) + flat_cache[value_offsets] = value_bytes + flat_cache[scale_offsets] = scale.view(torch.uint8).reshape(num_rows, scale_bytes) + + +def read_deepseek_v4_indexer_mxfp4_cache( + cache_2d: torch.Tensor, + slot_mapping: torch.Tensor, + block_size: int = 64, +) -> torch.Tensor: + """Dequantize MXFP4 indexer cache rows selected by `slot_mapping`.""" + + if cache_2d.dtype != torch.uint8: + raise TypeError(f"cache_2d must be uint8, got {cache_2d.dtype}") + index_head_dim, value_bytes, scale_bytes = _indexer_mxfp4_layout_from_cache( + cache_2d, block_size + ) + min_stride = block_size * (value_bytes + scale_bytes) + if cache_2d.dim() != 2 or cache_2d.shape[1] < min_stride: + raise ValueError( + f"cache_2d must be [pages, >= {min_stride}], got {tuple(cache_2d.shape)}" + ) + + out_shape = (slot_mapping.numel(), index_head_dim) + if slot_mapping.numel() == 0: + return torch.empty(out_shape, device=cache_2d.device, dtype=torch.float32) + + flat_cache = cache_2d.reshape(-1) + slots = slot_mapping.to(torch.int64) + valid = slots >= 0 + safe_slots = torch.where(valid, slots, torch.zeros_like(slots)) + pages = torch.div(safe_slots, block_size, rounding_mode="floor") + pos = safe_slots % block_size + page_base = pages * cache_2d.stride(0) + value_base = page_base + pos * value_bytes + scale_base = page_base + block_size * value_bytes + pos * scale_bytes + + value_offsets = ( + value_base[:, None] + + torch.arange( + value_bytes, + device=cache_2d.device, + dtype=torch.int64, + )[None, :] + ) + packed = flat_cache[value_offsets] + + scale_offsets = ( + scale_base[:, None] + + torch.arange( + scale_bytes, + device=cache_2d.device, + dtype=torch.int64, + )[None, :] + ) + scales = torch.pow(2.0, flat_cache[scale_offsets].to(torch.int32) - 127) + byte_scales = scales.float().repeat_interleave( + DEEPSEEK_V4_MXFP4_BLOCK_SIZE // 2, dim=1 + ) + + even = _e2m1_values(packed & 0xF) * byte_scales + odd = _e2m1_values(packed >> 4) * byte_scales + out = torch.empty(out_shape, device=cache_2d.device, dtype=torch.float32) + out[:, 0::2] = even + out[:, 1::2] = odd + return torch.where(valid[:, None], out, torch.zeros_like(out)) + + +def read_deepseek_v4_indexer_fp8_cache( + cache_2d: torch.Tensor, + slot_mapping: torch.Tensor, + block_size: int = 64, +) -> torch.Tensor: + """Dequantize FP8 indexer cache rows selected by `slot_mapping`.""" + + if cache_2d.dtype != torch.uint8: + raise TypeError(f"cache_2d must be uint8, got {cache_2d.dtype}") + index_head_dim, scale_bytes = _indexer_fp8_layout_from_cache(cache_2d, block_size) + min_stride = block_size * (index_head_dim + scale_bytes) + if cache_2d.dim() != 2 or cache_2d.shape[1] < min_stride: + raise ValueError( + f"cache_2d must be [pages, >= {min_stride}], got {tuple(cache_2d.shape)}" + ) + + out = torch.zeros( + slot_mapping.numel(), + index_head_dim, + device=cache_2d.device, + dtype=torch.float32, + ) + flat_cache = cache_2d.reshape(-1) + for token_idx, raw_slot in enumerate(slot_mapping.tolist()): + slot = int(raw_slot) + if slot < 0: + continue + page = slot // block_size + pos = slot % block_size + page_base = page * cache_2d.stride(0) + value_base = page_base + pos * index_head_dim + scale_base = page_base + block_size * index_head_dim + pos * scale_bytes + scale = flat_cache[scale_base : scale_base + scale_bytes].view(torch.float32)[0] + values = flat_cache[value_base : value_base + index_head_dim].view( + torch.float8_e4m3fn + ) + out[token_idx].copy_(values.float() * scale) + return out + + +def _compress_v4_state_windows_capturable( + state_cache: torch.Tensor, + token_to_req_indices: torch.Tensor, + positions: torch.Tensor, + compressor_slot_mapping: torch.Tensor, + block_table: torch.Tensor, + block_table_base_offsets: torch.Tensor | None, + compressor_block_size: int, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + compress_ratio: int, + head_dim: int, + overlap: bool, +) -> tuple[torch.Tensor, torch.Tensor]: + num_actual = min(compressor_slot_mapping.numel(), positions.numel()) + if num_actual == 0: + return ( + torch.empty((0, head_dim), device=state_cache.device, dtype=torch.float32), + torch.empty((0,), device=state_cache.device, dtype=torch.bool), + ) + + token_positions = positions[:num_actual].to(torch.int64) + state_slots = compressor_slot_mapping[:num_actual].to(torch.int64) + valid_token = (state_slots >= 0) & ( + torch.remainder(token_positions + 1, compress_ratio) == 0 + ) + + window = (2 if overlap else 1) * compress_ratio + offsets = torch.arange(window, device=state_cache.device, dtype=torch.int64) + window_positions = token_positions[:, None] - window + 1 + offsets[None, :] + table_idx_raw = torch.div( + window_positions, compressor_block_size, rounding_mode="floor" + ) + req_idx = token_to_req_indices[:num_actual].to(torch.int64).clamp_min(0) + if block_table_base_offsets is not None: + safe_req_for_base = req_idx.clamp( + 0, max(int(block_table_base_offsets.shape[0]) - 1, 0) + ) + base_logical_page = block_table_base_offsets.to( + device=state_cache.device, + dtype=torch.int64, + )[safe_req_for_base] + table_idx_raw = table_idx_raw - base_logical_page[:, None] + valid_window = ( + (window_positions >= 0) + & (table_idx_raw >= 0) + & (table_idx_raw < block_table.shape[1]) + ) + table_idx = table_idx_raw.clamp(0, max(block_table.shape[1] - 1, 0)) + block_number = block_table[req_idx[:, None], table_idx] + valid_window = valid_window & (block_number >= 0) + + safe_block = block_number.to(torch.int64).clamp_min(0) + pos_in_block = torch.remainder(window_positions.clamp_min(0), compressor_block_size) + rows = state_cache[safe_block, pos_in_block] + state_width = state_cache.shape[-1] // 2 + + if overlap: + head_offsets = torch.where( + offsets >= compress_ratio, + torch.full_like(offsets, head_dim), + torch.zeros_like(offsets), + ) + else: + head_offsets = torch.zeros_like(offsets) + dim_indices = ( + head_offsets[:, None] + + torch.arange(head_dim, device=state_cache.device, dtype=torch.int64)[None, :] + ) + dim_indices = dim_indices[None, :, :].expand(num_actual, -1, -1) + + kv_rows = torch.gather(rows[..., :state_width], -1, dim_indices).float() + score_rows = torch.gather(rows[..., state_width:], -1, dim_indices).float() + valid_window_f = valid_window.unsqueeze(-1) + score_rows = torch.where( + valid_window_f, score_rows, score_rows.new_full((), -1.0e30) + ) + weights = torch.softmax(score_rows, dim=1) + kv_rows = torch.where(valid_window_f, kv_rows, torch.zeros_like(kv_rows)) + compressed = torch.sum(kv_rows * weights, dim=1) + variance = compressed.square().sum(dim=-1, keepdim=True) / float(head_dim) + normed = compressed * torch.rsqrt(variance + rms_norm_eps) + return normed * rms_norm_weight.float(), valid_token + + +def deepseek_v4_hca_compress_kv_cache_insert( + state_cache: torch.Tensor, + token_to_req_indices: torch.Tensor, + positions: torch.Tensor, + compressor_slot_mapping: torch.Tensor, + block_table: torch.Tensor, + compressor_block_size: int, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + kv_cache_2d: torch.Tensor, + kv_slot_mapping: torch.Tensor, + kv_cache_block_size: int, + compress_ratio: int = 128, + block_table_base_offsets: torch.Tensor | None = None, +) -> None: + """Compress HCA state, normalize/RoPE/FP8-quantize, and insert KV cache. + + The HCA path writes one compressed cache entry only at positions where + `(position + 1) % 128 == 0`. + """ + + if compress_ratio != 128: + raise ValueError( + f"HCA cache insert requires compress_ratio=128, got {compress_ratio}" + ) + if state_cache.dim() != 3: + raise ValueError(f"state_cache must be 3D, got {tuple(state_cache.shape)}") + state_width = state_cache.shape[-1] // 2 + head_dim = int(rms_norm_weight.numel()) + if state_width != head_dim: + raise ValueError(f"HCA state width must be {head_dim}, got {state_width}") + if compressor_block_size != state_cache.shape[1]: + raise ValueError( + "compressor_block_size must match state_cache page size, " + f"got {compressor_block_size} vs {state_cache.shape[1]}" + ) + rope_dim = int(cos_sin_cache.shape[-1]) + min_block_stride = kv_cache_block_size * deepseek_v4_swa_row_bytes( + state_width, rope_dim + ) + if kv_cache_2d.dim() != 2 or kv_cache_2d.shape[1] < min_block_stride: + raise ValueError( + f"kv_cache_2d must be [blocks, >= {min_block_stride}] uint8, " + f"got {tuple(kv_cache_2d.shape)}" + ) + if kv_cache_2d.dtype != torch.uint8: + raise TypeError(f"kv_cache_2d must be uint8, got {kv_cache_2d.dtype}") + + num_actual = min( + compressor_slot_mapping.numel(), + positions.numel(), + kv_slot_mapping.numel(), + ) + if num_actual == 0: + return + if not state_cache.is_cuda: + raise ValueError( + "deepseek_v4_hca_compress_kv_cache_insert only supports CUDA tensors." + ) + + _triton_fused_sparse_compress_cache_insert( + state_cache=state_cache, + token_to_req_indices=token_to_req_indices, + positions=positions, + compressor_slot_mapping=compressor_slot_mapping, + block_table=block_table, + compressor_block_size=compressor_block_size, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + cos_sin_cache=cos_sin_cache, + kv_cache_2d=kv_cache_2d, + kv_slot_mapping=kv_slot_mapping, + kv_cache_block_size=kv_cache_block_size, + compress_ratio=compress_ratio, + overlap=False, + block_table_base_offsets=block_table_base_offsets, + ) + + +def deepseek_v4_csa_compress_kv_cache_insert( + state_cache: torch.Tensor, + token_to_req_indices: torch.Tensor, + positions: torch.Tensor, + compressor_slot_mapping: torch.Tensor, + block_table: torch.Tensor, + compressor_block_size: int, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + kv_cache_2d: torch.Tensor, + kv_slot_mapping: torch.Tensor, + kv_cache_block_size: int, + compress_ratio: int = 4, + block_table_base_offsets: torch.Tensor | None = None, +) -> None: + """Compress CSA state and insert one `fp8_ds_mla` row per 4 tokens. + + CSA uses overlap: the compression window spans eight token positions and + selects the first 512-wide slice from the older four positions and the + second slice from the newer four positions before the softmax-weighted sum. + """ + + if compress_ratio != 4: + raise ValueError( + f"CSA cache insert requires compress_ratio=4, got {compress_ratio}" + ) + if state_cache.dim() != 3: + raise ValueError(f"state_cache must be 3D, got {tuple(state_cache.shape)}") + state_width = state_cache.shape[-1] // 2 + head_dim = int(rms_norm_weight.numel()) + expected_width = head_dim * 2 + if state_width != expected_width: + raise ValueError(f"CSA state width must be {expected_width}, got {state_width}") + if compressor_block_size != state_cache.shape[1]: + raise ValueError( + "compressor_block_size must match state_cache page size, " + f"got {compressor_block_size} vs {state_cache.shape[1]}" + ) + rope_dim = int(cos_sin_cache.shape[-1]) + min_block_stride = kv_cache_block_size * deepseek_v4_swa_row_bytes( + head_dim, rope_dim + ) + if kv_cache_2d.dim() != 2 or kv_cache_2d.shape[1] < min_block_stride: + raise ValueError( + f"kv_cache_2d must be [blocks, >= {min_block_stride}] uint8, " + f"got {tuple(kv_cache_2d.shape)}" + ) + if kv_cache_2d.dtype != torch.uint8: + raise TypeError(f"kv_cache_2d must be uint8, got {kv_cache_2d.dtype}") + + num_actual = min(compressor_slot_mapping.numel(), positions.numel()) + if num_actual == 0: + return + if not state_cache.is_cuda: + raise ValueError( + "deepseek_v4_csa_compress_kv_cache_insert only supports CUDA tensors." + ) + + _triton_fused_sparse_compress_cache_insert( + state_cache=state_cache, + token_to_req_indices=token_to_req_indices, + positions=positions, + compressor_slot_mapping=compressor_slot_mapping, + block_table=block_table, + compressor_block_size=compressor_block_size, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + cos_sin_cache=cos_sin_cache, + kv_cache_2d=kv_cache_2d, + kv_slot_mapping=kv_slot_mapping, + kv_cache_block_size=kv_cache_block_size, + compress_ratio=compress_ratio, + overlap=True, + block_table_base_offsets=block_table_base_offsets, + ) + + +def deepseek_v4_csa_indexer_cache_insert( + state_cache: torch.Tensor, + token_to_req_indices: torch.Tensor, + positions: torch.Tensor, + compressor_slot_mapping: torch.Tensor, + block_table: torch.Tensor, + compressor_block_size: int, + rms_norm_weight: torch.Tensor, + rms_norm_eps: float, + cos_sin_cache: torch.Tensor, + kv_cache_2d: torch.Tensor, + kv_slot_mapping: torch.Tensor, + kv_cache_block_size: int, + use_fp4_cache: bool, + compress_ratio: int = 4, + block_table_base_offsets: torch.Tensor | None = None, +) -> None: + """Compress CSA indexer state and insert FP8/MXFP4 indexer cache rows.""" + + if compress_ratio != 4: + raise ValueError( + f"CSA indexer cache insert requires compress_ratio=4, got {compress_ratio}" + ) + if state_cache.dim() != 3: + raise ValueError(f"state_cache must be 3D, got {tuple(state_cache.shape)}") + state_width = state_cache.shape[-1] // 2 + index_head_dim = int(rms_norm_weight.numel()) + expected_width = index_head_dim * 2 + if state_width != expected_width: + raise ValueError( + f"CSA indexer state width must be {expected_width}, got {state_width}" + ) + + num_actual = min(compressor_slot_mapping.numel(), positions.numel()) + if num_actual == 0: + return + if not state_cache.is_cuda: + raise ValueError( + "deepseek_v4_csa_indexer_cache_insert only supports CUDA tensors." + ) + if use_fp4_cache: + _triton_fused_csa_indexer_mxfp4_cache_insert( + state_cache=state_cache, + token_to_req_indices=token_to_req_indices, + positions=positions, + compressor_slot_mapping=compressor_slot_mapping, + block_table=block_table, + compressor_block_size=compressor_block_size, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + cos_sin_cache=cos_sin_cache, + kv_cache_2d=kv_cache_2d, + kv_slot_mapping=kv_slot_mapping, + kv_cache_block_size=kv_cache_block_size, + compress_ratio=compress_ratio, + block_table_base_offsets=block_table_base_offsets, + ) + return + + normed, valid = _compress_v4_state_windows_capturable( + state_cache=state_cache, + token_to_req_indices=token_to_req_indices, + positions=positions, + compressor_slot_mapping=compressor_slot_mapping, + block_table=block_table, + block_table_base_offsets=block_table_base_offsets, + compressor_block_size=compressor_block_size, + rms_norm_weight=rms_norm_weight, + rms_norm_eps=rms_norm_eps, + compress_ratio=compress_ratio, + head_dim=index_head_dim, + overlap=True, + ) + compressed_positions = ( + torch.div( + positions[:num_actual].to(torch.int64), + compress_ratio, + rounding_mode="floor", + ) + * compress_ratio + ) + rotated = _apply_gptj_rope_tail_rows( + normed, + compressed_positions, + cos_sin_cache, + int(cos_sin_cache.shape[-1]), + ) + rotated = _deepseek_v4_hadamard_rotate(rotated).float() + _write_deepseek_v4_indexer_fp8_cache_capturable( + rotated, + kv_cache_2d, + kv_slot_mapping[:num_actual], + valid, + block_size=kv_cache_block_size, + ) diff --git a/python/tokenspeed/runtime/layers/attention/dsa/utils.py b/python/tokenspeed/runtime/layers/attention/dsa/utils.py new file mode 100644 index 0000000..815e177 --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/dsa/utils.py @@ -0,0 +1,52 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import torch + + +def workspace_indices_to_kv_slots( + workspace_indices: torch.Tensor, + kv_workspace_slots: torch.Tensor | None, +) -> torch.Tensor: + """Map DSA workspace-local top-k indices to global KV cache slot ids. + + Args: + workspace_indices: Top-k indices in the compact DSA prefill workspace. + Negative entries are treated as invalid sentinels and preserved. + kv_workspace_slots: Lookup table mapping workspace rows to KV cache slots. + + Returns: + A tensor with the same shape as ``workspace_indices`` containing int32 KV + cache slot ids, or int32 ``workspace_indices`` when no lookup is provided. + """ + if kv_workspace_slots is None or workspace_indices.numel() == 0: + return workspace_indices.to(torch.int32) + + flat_indices = workspace_indices.reshape(-1) + valid = flat_indices >= 0 + flat_slots = flat_indices.to(torch.int64) + if valid.any(): + flat_slots[valid] = kv_workspace_slots.to( + device=workspace_indices.device, + dtype=torch.int64, + ).index_select(0, flat_slots[valid]) + return flat_slots.view_as(workspace_indices).to(torch.int32) diff --git a/python/tokenspeed/runtime/layers/attention/kv_cache/base.py b/python/tokenspeed/runtime/layers/attention/kv_cache/base.py new file mode 100644 index 0000000..7e2f146 --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/kv_cache/base.py @@ -0,0 +1,159 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +from tokenspeed.runtime.configs.paged_cache_spec import PagedCacheGroupSpec +from tokenspeed.runtime.layers.paged_attention import PagedAttention +from tokenspeed.runtime.utils import get_colorful_logger + +if TYPE_CHECKING: + from tokenspeed.runtime.cache.kvstore_controller import LayerDoneCounter + +logger = get_colorful_logger(__name__) + + +class BaseTokenToKVPool: + """A memory pool that maps a token location to its kv cache data.""" + + paged_cache_group_specs: tuple[PagedCacheGroupSpec, ...] = () + paged_cache_group_page_counts: dict[str, int] = {} + supports_hierarchical_kv_cache: bool = True + + def __init__( + self, + size: int, + dtype: torch.dtype, + device: str, + max_batch_size: int, + max_context_len: int, + page_size: int, + rank: int, + ): + self.dtype = dtype + self.rank = rank + self.size = size + self.page_size = page_size + if dtype in (torch.float8_e5m2, torch.float8_e4m3fn): + # Store as torch.uint8 because Tensor.index_put is not implemented for torch.float8_e5m2 + self.store_dtype = torch.uint8 + else: + self.store_dtype = dtype + self.device = device + self.offload_chunk_page_num = 1024 + self.token_slot_refs = None + + # default state for optional layer-wise transfer control + self.layer_transfer_counter = None + logger.info( + f"Initialized token to kv pool with size {size}, dtype {dtype}, device {device}, page size {page_size}, rank {rank}" + ) + + @classmethod + def cell_size(self) -> int: + raise NotImplementedError() + + def register_layer_transfer_counter(self, layer_transfer_counter: LayerDoneCounter): + self.layer_transfer_counter = layer_transfer_counter + + def set_token_slot_refs(self, token_slot_refs: torch.Tensor): + self.token_slot_refs = token_slot_refs + + def bind_paged_cache_scheduler(self, scheduler: object) -> None: + """Optional hook for model-specific paged-cache diagnostics.""" + return None + + @torch.no_grad() + def clear_kv_buffers(self) -> None: + """Zero the KV buffers in place. + + Used by sleep/wake: after resume_memory_occupation re-maps the KV region + its pages hold garbage, so zero them. Subclasses store buffers under + different attributes (``k_buffer``/``v_buffer`` for MHA, ``kv_buffer`` — + possibly tuples — for MLA); introspect the known names so every pool is + covered without per-class overrides. For non-quantized KV this is + belt-and-suspenders (paging overwrites); for FP8 KV it removes garbage. + """ + attrs = ( + "k_buffer", + "v_buffer", + "kv_buffer", + # DeepSeek V4 pool buffer names. + "swa_kv_buffer", + "compressed_kv_buffer", + "compressor_state_buffer", + "indexer_kv_buffer", + "indexer_state_buffer", + ) + for attr in attrs: + for entry in getattr(self, attr, None) or []: + items = entry if isinstance(entry, (tuple, list)) else (entry,) + for t in items: + if torch.is_tensor(t): + t.zero_() + + def maybe_log_paged_cache_group_pages(self) -> None: + return None + + def get_key_buffer(self, layer_id: int) -> torch.Tensor: + raise NotImplementedError() + + def get_value_buffer(self, layer_id: int) -> torch.Tensor: + raise NotImplementedError() + + def get_kv_buffer(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: + raise NotImplementedError() + + def set_kv_buffer( + self, + layer: PagedAttention, + loc: torch.Tensor, + cache_k: torch.Tensor, + cache_v: torch.Tensor, + ) -> None: + raise NotImplementedError() + + def get_cpu_copy(self, page_indices: list[int]) -> torch.Tensor: + raise NotImplementedError() + + def load_cpu_copy( + self, kv_cache_cpu: torch.Tensor, page_indices: list[int] + ) -> None: + raise NotImplementedError() + + @property + def prefix_cache_required_group_ids(self) -> tuple[str, ...] | None: + """None means adjunct disabled; subclasses return required group ids.""" + return None + + # Buffer metadata used by prefill/decode disaggregation. + def get_contiguous_buf_infos(self): + raise NotImplementedError() + + def get_contiguous_buf_unit_lens(self): + return [1] * len(self.get_contiguous_buf_infos()[2]) + + # Layerwise buffer offsets used by prefill/decode disaggregation. + def get_layerwise_buf_info_offsets(self, start_idx=0): + raise NotImplementedError() diff --git a/python/tokenspeed/runtime/layers/attention/kv_cache/deepseek_v4.py b/python/tokenspeed/runtime/layers/attention/kv_cache/deepseek_v4.py new file mode 100644 index 0000000..dd47858 --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/kv_cache/deepseek_v4.py @@ -0,0 +1,1281 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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 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. + +from __future__ import annotations + +import logging +from collections.abc import Iterable, Sequence +from dataclasses import dataclass, field +from fractions import Fraction +from typing import Any + +import numpy as np +import torch + +from tokenspeed.runtime.configs.deepseek_v4_cache_spec import ( + DEEPSEEK_V4_COMPRESSED_LOGICAL_BLOCK_SIZE, + V4_INDEXER_COMPRESSOR_STATE_GROUP_ID, + V4_KERNEL_BLOCK_ROWS, + V4_SWA_KV_GROUP_ID, + build_v4_cache_specs, + deepseek_v4_indexer_fp8_row_bytes, + deepseek_v4_indexer_mxfp4_row_bytes, + deepseek_v4_swa_scale_dim, + deepseek_v4_swa_token_stride, + parse_v4_compressor_state_group_id, + v4_compressed_kv_group_id, + v4_compressor_state_group_id, +) +from tokenspeed.runtime.configs.paged_cache_spec import ( + PagedCacheGroupSpec, + compute_paged_cache_group_page_counts, +) +from tokenspeed.runtime.layers.attention.deepseek_v4_ops import ( + deepseek_v4_compressed_slot_mapping, +) +from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool +from tokenspeed.runtime.utils import get_colorful_logger +from tokenspeed.runtime.utils.common import ceil_div +from tokenspeed.runtime.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter + +logger = get_colorful_logger(__name__) + + +@dataclass(frozen=True) +class DeepseekV4CacheLayout: + layer_ratio: tuple[int, ...] + head_dim: int + rope_head_dim: int + page_size: int + use_fp4_indexer_cache: bool + index_head_dim: int = 128 + + @property + def swa_token_stride(self) -> int: + return deepseek_v4_swa_token_stride(self.head_dim, self.rope_head_dim) + + @property + def swa_scale_dim(self) -> int: + return deepseek_v4_swa_scale_dim(self.head_dim, self.rope_head_dim) + + @property + def swa_row_bytes(self) -> int: + return self.swa_token_stride + self.swa_scale_dim + + def swa_block_bytes(self, page_size: int | None = None) -> int: + if page_size is None: + page_size = self.page_size + block_bytes = page_size * self.swa_row_bytes + alignment = self.swa_token_stride + return ((block_bytes + alignment - 1) // alignment) * alignment + + def swa_cell_bytes(self) -> int: + block_bytes = self.swa_block_bytes() + return (block_bytes + self.page_size - 1) // self.page_size + + def storage_block_size(self, compress_ratio: int) -> int: + if compress_ratio > 1: + return max(1, DEEPSEEK_V4_COMPRESSED_LOGICAL_BLOCK_SIZE // compress_ratio) + return self.page_size + + def compressor_state_block_size(self, compress_ratio: int) -> int: + if compress_ratio == 4: + return 4 + if compress_ratio == 128: + return 8 + return self.page_size + + def compressed_cell_bytes(self, compress_ratio: int) -> int: + block_bytes = self.swa_block_bytes(self.storage_block_size(compress_ratio)) + return (block_bytes + self.page_size - 1) // self.page_size + + @property + def indexer_row_bytes(self) -> int: + if self.use_fp4_indexer_cache: + return deepseek_v4_indexer_mxfp4_row_bytes(self.index_head_dim) + return deepseek_v4_indexer_fp8_row_bytes(self.index_head_dim) + + def state_width(self, layer_id: int, *, indexer: bool = False) -> int: + if indexer: + return self.index_head_dim * 2 + return self.head_dim * (2 if self.layer_ratio[layer_id] == 4 else 1) + + def cache_cell_size(self, layer_num: int | None = None) -> int: + """Return bytes per token for the current V4 cache allocation layout.""" + if layer_num is None: + layer_num = len(self.layer_ratio) + if layer_num > len(self.layer_ratio): + raise ValueError( + "DeepSeek V4 cache layout has fewer layer ratios " + f"({len(self.layer_ratio)}) than requested layers ({layer_num})" + ) + + fp32_size = torch._utils._element_size(torch.float32) + cell_size = 0 + for layer_id in range(layer_num): + ratio = self.layer_ratio[layer_id] + cell_size += self.swa_cell_bytes() + if ratio > 1: + cell_size += self.compressed_cell_bytes(ratio) + cell_size += self.state_width(layer_id) * 2 * fp32_size + if ratio == 4: + indexer_block_bytes = ( + self.storage_block_size(ratio) * self.indexer_row_bytes + ) + cell_size += ( + indexer_block_bytes + self.page_size - 1 + ) // self.page_size + cell_size += self.state_width(layer_id, indexer=True) * 2 * fp32_size + return cell_size + + +def _deepseek_v4_cache_group_page_bytes( + layout: DeepseekV4CacheLayout, + specs: Sequence[PagedCacheGroupSpec], + layer_num: int, +) -> dict[str, int]: + if layer_num > len(layout.layer_ratio): + raise ValueError( + "DeepSeek V4 cache layout has fewer layer ratios " + f"({len(layout.layer_ratio)}) than requested layers ({layer_num})" + ) + + group_rows = {spec.group_id: int(spec.rows_per_page) for spec in specs} + page_bytes = {spec.group_id: 0 for spec in specs} + fp32_size = torch._utils._element_size(torch.float32) + + swa_block_bytes = layout.swa_block_bytes( + group_rows.get(V4_SWA_KV_GROUP_ID, V4_KERNEL_BLOCK_ROWS) + ) + for layer_id, ratio in enumerate(layout.layer_ratio[:layer_num]): + page_bytes[V4_SWA_KV_GROUP_ID] += swa_block_bytes + if ratio <= 1: + continue + + compressed_group_id = v4_compressed_kv_group_id(ratio) + compressed_block_size = layout.storage_block_size(ratio) + page_bytes[compressed_group_id] += layout.swa_block_bytes(compressed_block_size) + + state_group_id = v4_compressor_state_group_id(ratio) + state_block_size = group_rows.get(state_group_id, layout.page_size) + page_bytes[state_group_id] += ( + state_block_size * layout.state_width(layer_id) * 2 * fp32_size + ) + + if ratio == 4: + indexer_block_size = max(V4_KERNEL_BLOCK_ROWS, compressed_block_size) + page_bytes[compressed_group_id] += ( + indexer_block_size * layout.indexer_row_bytes + ) + + indexer_state_block_size = group_rows.get( + V4_INDEXER_COMPRESSOR_STATE_GROUP_ID, + layout.compressor_state_block_size(ratio), + ) + page_bytes[V4_INDEXER_COMPRESSOR_STATE_GROUP_ID] += ( + indexer_state_block_size + * layout.state_width(layer_id, indexer=True) + * 2 + * fp32_size + ) + + return page_bytes + + +def _estimate_deepseek_v4_cache_bytes( + *, + layout: DeepseekV4CacheLayout, + hf_config: Any, + layer_num: int, + max_total_tokens: int, + max_live_requests: int, + max_scheduled_tokens: int, + max_context_len: int, +) -> int: + """Estimate bytes allocated by DeepseekV4TokenToKVPool for a token budget.""" + if layer_num > len(layout.layer_ratio): + raise ValueError( + "DeepSeek V4 cache layout has fewer layer ratios " + f"({len(layout.layer_ratio)}) than requested layers ({layer_num})" + ) + if max_total_tokens < 0: + raise ValueError(f"max_total_tokens must be >= 0, got {max_total_tokens}") + + specs = tuple(build_v4_cache_specs(hf_config, layer_ratio=layout.layer_ratio)) + page_bytes = _deepseek_v4_cache_group_page_bytes(layout, specs, layer_num) + counts = compute_paged_cache_group_page_counts( + specs, + max_live_requests=max_live_requests, + max_scheduled_tokens=max(0, int(max_scheduled_tokens)), + max_total_tokens=max_total_tokens, + max_context_len=max_context_len, + ) + return int( + sum( + int(counts[gid]) * bytes_per_page + for gid, bytes_per_page in page_bytes.items() + ) + ) + + +def profile_deepseek_v4_max_num_pages( + *, + layout: DeepseekV4CacheLayout, + hf_config: Any, + layer_num: int, + max_live_requests: int, + max_scheduled_tokens: int, + max_context_len: int, + available_cache_memory_bytes: int, + draft_cache_cell_size: int = 0, + decode_input_tokens: int = 1, + overlap_schedule_depth: int = 0, +) -> int: + """Return the largest scheduler page budget that fits V4 grouped caches.""" + page_size = int(layout.page_size) + if page_size <= 0: + raise ValueError(f"page_size must be positive, got {page_size}") + if available_cache_memory_bytes <= 0: + return 0 + if draft_cache_cell_size < 0: + raise ValueError( + f"draft_cache_cell_size must be >= 0, got {draft_cache_cell_size}" + ) + + draft_cache_cell_size = int(draft_cache_cell_size) + max_live_requests = int(max_live_requests) + max_scheduled_tokens = max(0, int(max_scheduled_tokens)) + max_context_len = int(max_context_len) + specs = tuple(build_v4_cache_specs(hf_config, layer_ratio=layout.layer_ratio)) + page_bytes = _deepseek_v4_cache_group_page_bytes(layout, specs, layer_num) + + def _bytes_for_pages(num_pages: int) -> int: + num_tokens = int(num_pages) * page_size + counts = compute_paged_cache_group_page_counts( + specs, + max_live_requests=max_live_requests, + max_scheduled_tokens=max_scheduled_tokens, + max_total_tokens=num_tokens, + max_context_len=max_context_len, + decode_input_tokens=decode_input_tokens, + overlap_schedule_depth=overlap_schedule_depth, + ) + cache_bytes = sum( + int(counts[gid]) * bytes_per_page + for gid, bytes_per_page in page_bytes.items() + ) + return int(cache_bytes + num_tokens * draft_cache_cell_size) + + if _bytes_for_pages(1) > available_cache_memory_bytes: + return 0 + + if not any(int(ratio) > 1 for ratio in layout.layer_ratio[:layer_num]): + return max( + 1, + (int(max_live_requests) * int(max_context_len) + page_size - 1) + // page_size, + ) + + # Fixed bytes cover resident sliding windows, request fragments, and dummy + # pages. Variable bytes are piecewise linear before and after the global + # scheduled-token write budget is capped. + fixed_counts = compute_paged_cache_group_page_counts( + specs, + max_live_requests=max_live_requests, + max_scheduled_tokens=max_scheduled_tokens, + max_total_tokens=0, + max_context_len=max_context_len, + decode_input_tokens=decode_input_tokens, + overlap_schedule_depth=overlap_schedule_depth, + ) + fixed_bytes = sum( + int(fixed_counts[gid]) * bytes_per_page + for gid, bytes_per_page in page_bytes.items() + ) + full_history_slope = Fraction(page_size * draft_cache_cell_size, 1) + scheduled_slope = Fraction(0, 1) + scheduled_cap_bytes = 0 + for spec in specs: + bytes_per_page = page_bytes[spec.group_id] + if bytes_per_page == 0: + continue + raw_per_page = int(spec.rows_per_page) * int(spec.entry_stride_tokens) + if spec.retention == "full_history": + full_history_slope += Fraction(page_size * bytes_per_page, raw_per_page) + elif spec.retention == "sliding_window": + scheduled_slope += Fraction(page_size * bytes_per_page, raw_per_page) + scheduled_cap_bytes += ( + ceil_div(max_scheduled_tokens, raw_per_page) * bytes_per_page + ) + + def _pages_from_budget(extra_bytes: int, slope: Fraction) -> int: + if extra_bytes <= 0 or slope <= 0: + return 0 + return int(extra_bytes * slope.denominator // slope.numerator) + + cap_pages = ceil_div(max_scheduled_tokens, page_size) + candidate = 0 + pre_cap_slope = full_history_slope + scheduled_slope + if cap_pages > 0: + pre_cap_pages = _pages_from_budget( + available_cache_memory_bytes - fixed_bytes, + pre_cap_slope, + ) + candidate = min(pre_cap_pages, cap_pages - 1) + + post_cap_fixed_bytes = fixed_bytes + scheduled_cap_bytes + post_cap_pages = _pages_from_budget( + available_cache_memory_bytes - post_cap_fixed_bytes, + full_history_slope, + ) + if post_cap_pages >= cap_pages: + candidate = max(candidate, post_cap_pages) + candidate = max(1, candidate) + + while candidate > 0 and _bytes_for_pages(candidate) > available_cache_memory_bytes: + candidate -= 1 + while _bytes_for_pages(candidate + 1) <= available_cache_memory_bytes: + candidate += 1 + return int(candidate) + + +def _split_paged_cache_block_tables_into_v4_metadata( + paged_cache_block_tables: dict[str, torch.Tensor], + paged_cache_block_table_base_offsets: dict[str, torch.Tensor] | None = None, +) -> tuple[ + torch.Tensor | None, + dict[int, torch.Tensor], + torch.Tensor | None, + torch.Tensor | None, + dict[int, torch.Tensor], + torch.Tensor | None, +]: + """Split paged-cache dict into V4-named tables + per-sliding-group offsets. + + Returns (swa, {ratio: compressor_state}, indexer_state, swa_base, + {ratio: compressor_state_base}, indexer_state_base). Unknown group ids + are ignored. Base offsets are None / missing when the input lacks them. + """ + offsets = paged_cache_block_table_base_offsets or {} + swa = paged_cache_block_tables.get(V4_SWA_KV_GROUP_ID) + indexer_state = paged_cache_block_tables.get(V4_INDEXER_COMPRESSOR_STATE_GROUP_ID) + swa_base = offsets.get(V4_SWA_KV_GROUP_ID) + indexer_state_base = offsets.get(V4_INDEXER_COMPRESSOR_STATE_GROUP_ID) + compressor_state: dict[int, torch.Tensor] = {} + compressor_state_base: dict[int, torch.Tensor] = {} + for gid, table in paged_cache_block_tables.items(): + ratio = parse_v4_compressor_state_group_id(gid) + if ratio is None: + continue + compressor_state[ratio] = table + base = offsets.get(gid) + if base is not None: + compressor_state_base[ratio] = base + return ( + swa, + compressor_state, + indexer_state, + swa_base, + compressor_state_base, + indexer_state_base, + ) + + +def _safe_page_ids( + block_table: torch.Tensor, + req_indices: torch.Tensor, + page_indices: torch.Tensor, +) -> torch.Tensor: + req_i64 = req_indices.to(torch.int64) + page_i64 = page_indices.to(torch.int64) + sentinel = torch.full_like(page_i64, -1, dtype=torch.int64) + rows = int(block_table.shape[0]) if block_table.ndim >= 1 else 0 + cols = int(block_table.shape[1]) if block_table.ndim >= 2 else 0 + if rows <= 0 or cols <= 0: + return sentinel + valid = (req_i64 >= 0) & (req_i64 < rows) & (page_i64 >= 0) & (page_i64 < cols) + safe_req = req_i64.clamp(0, rows - 1) + safe_page = page_i64.clamp(0, cols - 1) + page_ids = block_table[safe_req, safe_page].to(torch.int64) + return torch.where(valid, page_ids, sentinel) + + +def _expand_group_values_for_tokens( + values: torch.Tensor, + num_tokens: int, + name: str, +) -> torch.Tensor: + if values.numel() == num_tokens: + return values + if values.numel() <= 0 or num_tokens % values.numel() != 0: + raise RuntimeError( + f"DeepSeek V4 {name} has incompatible shape for packed tokens: " + f"{values.numel()} entries for {num_tokens} tokens" + ) + return values.repeat_interleave(num_tokens // values.numel()) + + +def _group_slot_mapping_from_raw( + positions: torch.Tensor, + req_indices: torch.Tensor, + block_table: torch.Tensor, + rows_per_page: int, + entry_stride_tokens: int = 1, + base_offsets: torch.Tensor | None = None, +) -> torch.Tensor: + if rows_per_page <= 0: + raise ValueError(f"rows_per_page must be > 0, got {rows_per_page}") + if entry_stride_tokens <= 0: + raise ValueError(f"entry_stride_tokens must be > 0, got {entry_stride_tokens}") + pos_i64 = positions.to(torch.int64) + logical_row = torch.div(pos_i64, entry_stride_tokens, rounding_mode="floor") + logical_page = torch.div(logical_row, rows_per_page, rounding_mode="floor") + offsets = logical_row % rows_per_page + req_indices = _expand_group_values_for_tokens( + req_indices, + positions.numel(), + "request indices", + ) + table_page = logical_page + if base_offsets is not None: + req_i64 = req_indices.to(torch.int64) + rows = int(base_offsets.shape[0]) + if rows <= 0: + table_page = logical_page.new_full(logical_page.shape, -1) + else: + valid_req = (req_i64 >= 0) & (req_i64 < rows) + safe_req = req_i64.clamp(0, rows - 1) + base = base_offsets.to( + device=logical_page.device, + dtype=torch.int64, + )[safe_req] + table_page = torch.where(valid_req, logical_page - base, -1) + page_ids = _safe_page_ids(block_table, req_indices, table_page) + slots = page_ids * rows_per_page + offsets + return torch.where(page_ids >= 0, slots, torch.full_like(slots, -1)) + + +def _mask_invalid_graph_tokens( + slot_mapping: torch.Tensor, + is_valid_token: torch.Tensor | None, +) -> torch.Tensor: + if is_valid_token is None: + return slot_mapping + valid = _expand_group_values_for_tokens( + is_valid_token, + slot_mapping.numel(), + "slot validity mask", + ).to( + device=slot_mapping.device, + dtype=torch.bool, + ) + return torch.where(valid, slot_mapping, torch.full_like(slot_mapping, -1)) + + +def _compressed_boundary_mask( + positions: torch.Tensor, + compress_ratio: int, +) -> torch.Tensor: + if compress_ratio <= 1: + return torch.ones_like(positions, dtype=torch.bool) + return ((positions.to(torch.int64) + 1) % compress_ratio) == 0 + + +@dataclass +class DeepseekV4CacheMetadata: + page_size: int + block_table: torch.Tensor + paged_cache_block_tables: dict[str, torch.Tensor] = field(default_factory=dict) + # Per-sliding-group [num_reqs] int32 base logical-page offset that + # accompanies each compact block table. Consumers index sliding tables as + # logical_page - base_offset; full-history groups omit the key (base 0). + paged_cache_block_table_base_offsets: dict[str, torch.Tensor] = field( + default_factory=dict + ) + swa_block_table: torch.Tensor | None = None + swa_base_logical_page: torch.Tensor | None = None + compressor_state_block_tables: dict[int, torch.Tensor] = field(default_factory=dict) + compressor_state_base_logical_pages: dict[int, torch.Tensor] = field( + default_factory=dict + ) + indexer_state_block_table: torch.Tensor | None = None + indexer_state_base_logical_page: torch.Tensor | None = None + decode_compressed_slot_mappings: dict[tuple[int, int], torch.Tensor] = field( + default_factory=dict + ) + + def compressed_block_table( + self, + compress_ratio: int, + kv_cache_block_size: int | None = None, + ) -> torch.Tensor: + del kv_cache_block_size + if compress_ratio <= 1: + return self.block_table + table = self.paged_cache_block_tables.get( + v4_compressed_kv_group_id(compress_ratio) + ) + if table is None: + raise RuntimeError( + "DeepSeek V4 missing paged-cache block table for compressed " + f"KV group {v4_compressed_kv_group_id(compress_ratio)!r}" + ) + return table + + @staticmethod + def safe_page_ids( + block_table: torch.Tensor, + req_indices: torch.Tensor, + page_indices: torch.Tensor, + ) -> torch.Tensor: + return _safe_page_ids(block_table, req_indices, page_indices) + + def _update_decode_compressed_slot_mapping( + self, + *, + token_to_req_indices: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + compress_ratio: int, + kv_cache_block_size: int, + is_valid_token: torch.Tensor | None = None, + ) -> torch.Tensor: + num_tokens = token_to_req_indices.shape[0] + key = (compress_ratio, kv_cache_block_size) + out = self.decode_compressed_slot_mappings.get(key) + if out is None or out.shape[0] < num_tokens or out.device != seq_lens.device: + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "DeepSeek V4 compressed slot metadata must be allocated before " + "CUDA graph capture" + ) + with torch.inference_mode(False): + out = torch.empty(num_tokens, dtype=torch.int64, device=seq_lens.device) + self.decode_compressed_slot_mappings[key] = out + + block_table = self.compressed_block_table(compress_ratio, kv_cache_block_size) + if block_table is not self.block_table: + req_idx = token_to_req_indices[:num_tokens].to(torch.int64) + query_starts = query_start_loc[req_idx].to(torch.int64) + query_lens = query_start_loc[req_idx + 1].to(torch.int64) - query_starts + seq_lens_for_token = seq_lens[req_idx].to(torch.int64) + token_offsets = torch.arange( + num_tokens, + dtype=torch.int64, + device=seq_lens.device, + ) + positions = seq_lens_for_token - query_lens + token_offsets - query_starts + compressed_pos = torch.div( + positions, + compress_ratio, + rounding_mode="floor", + ) + page_indices = torch.div( + compressed_pos, + kv_cache_block_size, + rounding_mode="floor", + ) + offsets = compressed_pos % kv_cache_block_size + base_offsets = self.paged_cache_block_table_base_offsets.get( + v4_compressed_kv_group_id(compress_ratio) + ) + if base_offsets is not None: + page_indices = ( + page_indices + - base_offsets.to( + device=page_indices.device, + dtype=torch.int64, + )[req_idx] + ) + page_ids = _safe_page_ids(block_table, req_idx, page_indices) + valid_slots = (page_ids >= 0) & _compressed_boundary_mask( + positions, + compress_ratio, + ) + slot_mapping = torch.where( + valid_slots, + page_ids * kv_cache_block_size + offsets, + torch.full_like(page_ids, -1), + ) + out.copy_(_mask_invalid_graph_tokens(slot_mapping, is_valid_token)) + return out + + mapping = deepseek_v4_compressed_slot_mapping( + num_tokens=num_tokens, + query_start_loc=query_start_loc, + seq_lens=seq_lens, + block_table=self.block_table, + block_size=kv_cache_block_size, + compress_ratio=compress_ratio, + out=out, + ) + if is_valid_token is not None: + mapping.copy_(_mask_invalid_graph_tokens(mapping, is_valid_token)) + return mapping + + def refresh_decode_compressed_slot_mappings( + self, + *, + token_to_req_indices: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + is_valid_token: torch.Tensor | None = None, + ) -> None: + for compress_ratio, kv_cache_block_size in list( + self.decode_compressed_slot_mappings + ): + self._update_decode_compressed_slot_mapping( + token_to_req_indices=token_to_req_indices, + query_start_loc=query_start_loc, + seq_lens=seq_lens, + compress_ratio=compress_ratio, + kv_cache_block_size=kv_cache_block_size, + is_valid_token=is_valid_token, + ) + + def compressed_slot_mapping( + self, + positions: torch.Tensor, + compress_ratio: int, + *, + token_to_req_indices: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens: torch.Tensor, + kv_cache_block_size: int | None = None, + use_decode_cache: bool = False, + is_valid_token: torch.Tensor | None = None, + ) -> torch.Tensor: + if kv_cache_block_size is None: + kv_cache_block_size = self.page_size + block_table = self.compressed_block_table(compress_ratio, kv_cache_block_size) + if ( + use_decode_cache + and positions.is_cuda + and (block_table.is_cuda or self.block_table.is_cuda) + ): + cached = self.decode_compressed_slot_mappings.get( + (compress_ratio, kv_cache_block_size) + ) + if ( + cached is not None + and cached.shape[0] >= positions.numel() + and cached.device == seq_lens.device + ): + return cached[: positions.numel()] + mapping = self._update_decode_compressed_slot_mapping( + token_to_req_indices=token_to_req_indices, + query_start_loc=query_start_loc, + seq_lens=seq_lens, + compress_ratio=compress_ratio, + kv_cache_block_size=kv_cache_block_size, + is_valid_token=is_valid_token, + ) + return mapping[: positions.numel()] + compressed_pos = torch.div( + positions.to(torch.int64), compress_ratio, rounding_mode="floor" + ) + page_indices = torch.div( + compressed_pos, kv_cache_block_size, rounding_mode="floor" + ) + offsets = compressed_pos % kv_cache_block_size + req_idx = token_to_req_indices[: positions.numel()].long() + if block_table is self.block_table: + page_ids = block_table[req_idx, page_indices.long()].to(torch.int64) + else: + base_offsets = self.paged_cache_block_table_base_offsets.get( + v4_compressed_kv_group_id(compress_ratio) + ) + if base_offsets is not None: + page_indices = ( + page_indices + - base_offsets.to( + device=page_indices.device, + dtype=torch.int64, + )[req_idx] + ) + page_ids = _safe_page_ids(block_table, req_idx, page_indices.long()) + slots = page_ids.to(torch.int64) * kv_cache_block_size + offsets + valid_slots = (page_ids >= 0) & _compressed_boundary_mask( + positions, + compress_ratio, + ) + slot_mapping = torch.where( + valid_slots, + slots, + torch.full_like(slots, -1), + ) + return _mask_invalid_graph_tokens(slot_mapping, is_valid_token) + + +def deepseek_v4_cache_layout_from_config( + hf_config, + page_size: int, + use_fp4_indexer_cache: bool, + layer_indices: Iterable[int] | None = None, +) -> DeepseekV4CacheLayout: + compress_ratios = tuple(hf_config.compress_ratios) + if layer_indices is None: + layer_ratios = compress_ratios + else: + layer_indices = tuple(layer_indices) + if any(idx < 0 or idx >= len(compress_ratios) for idx in layer_indices): + raise ValueError( + "DeepSeek V4 cache layout layer index out of range: " + f"indices={layer_indices}, ratios={len(compress_ratios)}" + ) + layer_ratios = [compress_ratios[idx] for idx in layer_indices] + raw_layer_ratios = tuple(int(x) for x in layer_ratios) + for ratio in raw_layer_ratios: + if ratio not in (0, 1, 4, 128): + raise ValueError( + "Unsupported DeepSeek V4 cache compress_ratio=" + f"{ratio}; expected one of 0, 1, 4, or 128" + ) + + return DeepseekV4CacheLayout( + layer_ratio=tuple(max(1, ratio) for ratio in raw_layer_ratios), + head_dim=int(hf_config.head_dim), + rope_head_dim=int(hf_config.qk_rope_head_dim), + page_size=page_size, + use_fp4_indexer_cache=use_fp4_indexer_cache, + index_head_dim=int(getattr(hf_config, "index_head_dim", 128)), + ) + + +class DeepseekV4TokenToKVPool(BaseTokenToKVPool): + """DeepSeek V4 fp8_ds_mla cache pool. + + TokenSpeed keeps SWA, compressed, compressor-state, and CSA indexer caches + in dedicated per-group paged pools (see PagedCacheGroup* on the scheduler + side and ``build_v4_cache_specs`` here), keeping ordinary MLA models on + their existing single-pool contract. The ``indexer_kv_buffer`` shares its + page table and page-count budget with the ``v4.c{ratio}a.compressed_kv`` + group rather than owning a separate group of its own. + """ + + supports_hierarchical_kv_cache = False + + def __init__( + self, + size: int, + model_dtype: torch.dtype, + layout: DeepseekV4CacheLayout, + layer_num: int, + device: str, + enable_memory_saver: bool, + max_batch_size: int, + max_context_len: int, + page_size: int, + rank: int, + hf_config: Any, + max_scheduled_tokens: int, + decode_input_tokens: int = 1, + overlap_schedule_depth: int = 0, + ) -> None: + if size <= 0: + raise ValueError(f"DeepSeek V4 KV pool size must be positive, got {size}") + if layer_num != len(layout.layer_ratio): + raise ValueError( + "DeepSeek V4 KV pool layer_num must match cache layout ratios: " + f"layer_num={layer_num}, ratios={len(layout.layer_ratio)}" + ) + super().__init__( + size=size, + dtype=torch.uint8, + device=device, + max_batch_size=max_batch_size, + max_context_len=max_context_len, + page_size=page_size, + rank=rank, + ) + # Tag KV allocations as "kv_cache" (no CPU backup: discarded on sleep) + # so release/resume_memory_occupation frees them. See memory_occupation.py. + self.memory_saver_adapter = TorchMemorySaverAdapter.create( + enable=enable_memory_saver + ) + self.model_dtype = model_dtype + self.layout = layout + self.layer_num = layer_num + self.max_batch_size = max_batch_size + self.max_context_len = max_context_len + self.num_pages = (size + page_size - 1) // page_size + 1 + self.paged_cache_group_specs = tuple( + build_v4_cache_specs(hf_config, layer_ratio=layout.layer_ratio) + ) + self._paged_cache_group_specs_by_id = { + spec.group_id: spec for spec in self.paged_cache_group_specs + } + self._paged_cache_scheduler: object | None = None + self._paged_cache_state_group_ids = tuple( + str(spec.group_id) + for spec in self.paged_cache_group_specs + if spec.family == "state" + ) + self.paged_cache_group_page_counts = compute_paged_cache_group_page_counts( + self.paged_cache_group_specs, + max_live_requests=max_batch_size, + max_scheduled_tokens=max(0, int(max_scheduled_tokens)), + max_total_tokens=size, + max_context_len=max_context_len, + decode_input_tokens=decode_input_tokens, + overlap_schedule_depth=overlap_schedule_depth, + ) + + def _group_rows(group_id: str, default: int) -> int: + spec = self._paged_cache_group_specs_by_id.get(group_id) + return int(spec.rows_per_page) if spec is not None else int(default) + + self.swa_block_size = _group_rows(V4_SWA_KV_GROUP_ID, V4_KERNEL_BLOCK_ROWS) + self.state_block_size = page_size + self.swa_block_bytes = layout.swa_block_bytes(self.swa_block_size) + self.compressed_block_sizes = tuple( + layout.storage_block_size(ratio) if ratio > 1 else page_size + for ratio in layout.layer_ratio + ) + self.indexer_block_sizes = tuple( + ( + max(V4_KERNEL_BLOCK_ROWS, self.compressed_block_sizes[layer_id]) + if ratio == 4 + else 0 + ) + for layer_id, ratio in enumerate(layout.layer_ratio) + ) + self.compressor_state_block_sizes = tuple( + ( + _group_rows(v4_compressor_state_group_id(ratio), page_size) + if ratio > 1 + else page_size + ) + for ratio in layout.layer_ratio + ) + self.indexer_state_block_sizes = tuple( + ( + _group_rows( + V4_INDEXER_COMPRESSOR_STATE_GROUP_ID, + layout.compressor_state_block_size(ratio), + ) + if ratio == 4 + else 0 + ) + for ratio in layout.layer_ratio + ) + self.compressed_block_size = ( + self.compressed_block_sizes[0] if self.compressed_block_sizes else page_size + ) + + swa_pages = self.paged_cache_group_page_counts.get( + V4_SWA_KV_GROUP_ID, + self.num_pages, + ) + with self.memory_saver_adapter.region(tag="kv_cache", enable_cpu_backup=False): + self.swa_kv_buffer = [ + torch.zeros( + (swa_pages, self.swa_block_bytes), + dtype=torch.uint8, + device=device, + ) + for _ in range(layer_num) + ] + self.compressed_kv_buffer: list[torch.Tensor | None] = [] + self.compressor_state_buffer: list[torch.Tensor | None] = [] + self.indexer_kv_buffer: list[torch.Tensor | None] = [] + self.indexer_state_buffer: list[torch.Tensor | None] = [] + for layer_id, ratio in enumerate(layout.layer_ratio): + has_compressed = ratio > 1 + has_indexer = ratio == 4 + compressed_block_size = self.compressed_block_sizes[layer_id] + compressed_group_id = v4_compressed_kv_group_id(ratio) + compressed_pages = self.num_pages + if has_compressed: + compressed_pages = self.paged_cache_group_page_counts.get( + compressed_group_id, + self.num_pages, + ) + self.compressed_kv_buffer.append( + torch.zeros( + ( + compressed_pages, + layout.swa_block_bytes(compressed_block_size), + ), + dtype=torch.uint8, + device=device, + ) + if has_compressed + else None + ) + compressor_state_block_size = self.compressor_state_block_sizes[ + layer_id + ] + compressor_state_group_id = v4_compressor_state_group_id(ratio) + compressor_state_pages = self.num_pages + if has_compressed: + compressor_state_pages = self.paged_cache_group_page_counts.get( + compressor_state_group_id, + self.num_pages, + ) + self.compressor_state_buffer.append( + torch.empty( + ( + compressor_state_pages, + compressor_state_block_size, + layout.state_width(layer_id) * 2, + ), + dtype=torch.float32, + device=device, + ) + if has_compressed + else None + ) + indexer_block_size = self.indexer_block_sizes[layer_id] + self.indexer_kv_buffer.append( + torch.zeros( + ( + compressed_pages, + indexer_block_size * layout.indexer_row_bytes, + ), + dtype=torch.uint8, + device=device, + ) + if has_indexer + else None + ) + indexer_state_block_size = self.indexer_state_block_sizes[layer_id] + indexer_state_pages = self.num_pages + if has_indexer: + indexer_state_pages = self.paged_cache_group_page_counts.get( + V4_INDEXER_COMPRESSOR_STATE_GROUP_ID, + self.num_pages, + ) + self.indexer_state_buffer.append( + torch.empty( + ( + indexer_state_pages, + indexer_state_block_size, + layout.state_width(layer_id, indexer=True) * 2, + ), + dtype=torch.float32, + device=device, + ) + if has_indexer + else None + ) + + logger.info( + "Initialized DeepSeek V4 KV pool: %d pages, %d layers, fp4 indexer=%s, compressed block sizes=%s", + self.num_pages, + layer_num, + layout.use_fp4_indexer_cache, + self.compressed_block_sizes, + ) + + @property + def prefix_cache_required_group_ids(self) -> tuple[str, ...]: + return tuple( + str(spec.group_id) + for spec in self.paged_cache_group_specs + if spec.family == "history" + ) + + def bind_paged_cache_scheduler(self, scheduler: object) -> None: + self._paged_cache_scheduler = scheduler + + def maybe_log_paged_cache_group_pages(self) -> None: + scheduler = self._paged_cache_scheduler + if self.rank != 0 or scheduler is None or not self._paged_cache_state_group_ids: + return + if not logger.isEnabledFor(logging.DEBUG): + return + + parts = [] + for group_id in self._paged_cache_state_group_ids: + total = scheduler.paged_cache_group_total_pages(group_id) + available = scheduler.paged_cache_group_available_pages(group_id) + failed = scheduler.paged_cache_group_failed_alloc_count(group_id) + parts.append( + f"{group_id}: used={total - available}/{total}, " + f"available={available}, failed_alloc={failed}" + ) + logger.debug("DeepSeek V4 paged-cache state group pages. %s", "; ".join(parts)) + + def _require( + self, buffers: list[torch.Tensor | None], layer_id: int, name: str + ) -> torch.Tensor: + buf = buffers[layer_id] + if buf is None: + raise ValueError(f"DeepSeek V4 layer {layer_id} has no {name} cache") + return buf + + def get_swa_kv_buffer(self, layer_id: int) -> torch.Tensor: + return self.swa_kv_buffer[layer_id] + + @property + def swa_capacity_slots(self) -> int: + """Writable SWA cache capacity shared by every layer, in token slots. + + Every layer's SWA buffer is allocated with the same page count, so a + single capacity (pages * tokens per block) bounds the write-slot + mapping shared across layers. Returns 0 when no SWA buffers exist; + callers must then mask all slots rather than skip the bounds check. + """ + if not self.swa_kv_buffer: + return 0 + return int(self.swa_kv_buffer[0].shape[0]) * int(self.swa_block_size) + + def get_compressed_kv_buffer_2d(self, layer_id: int) -> torch.Tensor: + return self._require(self.compressed_kv_buffer, layer_id, "compressed KV") + + def get_compressed_block_size(self, layer_id: int) -> int: + return self.compressed_block_sizes[layer_id] + + def get_indexer_block_size(self, layer_id: int) -> int: + block_size = self.indexer_block_sizes[layer_id] + if block_size <= 0: + raise ValueError(f"DeepSeek V4 layer {layer_id} has no indexer cache") + return block_size + + def get_compressor_state_block_size(self, layer_id: int) -> int: + block_size = self.compressor_state_block_sizes[layer_id] + if block_size <= 0: + raise ValueError( + f"DeepSeek V4 layer {layer_id} has no compressor state cache" + ) + return block_size + + def get_compressor_state_buffer(self, layer_id: int) -> torch.Tensor: + return self._require(self.compressor_state_buffer, layer_id, "compressor state") + + def get_compressor_state_view(self, layer_id: int) -> torch.Tensor: + buf = self.get_compressor_state_buffer(layer_id) + block_size = self.get_compressor_state_block_size(layer_id) + return buf.view(-1, block_size, buf.shape[-1]) + + def get_indexer_kv_buffer_2d(self, layer_id: int) -> torch.Tensor: + return self._require(self.indexer_kv_buffer, layer_id, "indexer KV") + + def get_indexer_state_block_size(self, layer_id: int) -> int: + block_size = self.indexer_state_block_sizes[layer_id] + if block_size <= 0: + raise ValueError(f"DeepSeek V4 layer {layer_id} has no indexer state cache") + return block_size + + def get_indexer_state_buffer(self, layer_id: int) -> torch.Tensor: + return self._require(self.indexer_state_buffer, layer_id, "indexer state") + + def get_indexer_state_view(self, layer_id: int) -> torch.Tensor: + buf = self.get_indexer_state_buffer(layer_id) + block_size = self.get_indexer_state_block_size(layer_id) + return buf.view(-1, block_size, buf.shape[-1]) + + def get_key_buffer(self, layer_id: int) -> torch.Tensor: + return self.get_swa_kv_buffer(layer_id) + + def get_value_buffer(self, layer_id: int) -> torch.Tensor: + return self.get_swa_kv_buffer(layer_id) + + def get_kv_buffer(self, layer_id: int): + buf = self.get_swa_kv_buffer(layer_id) + return buf, buf + + def set_kv_buffer(self, *args, **kwargs) -> None: + raise NotImplementedError( + "DeepSeek V4 writes KV cache through V4 attention helpers" + ) + + def _move_fp8_ds_mla_rows( + self, + buf: torch.Tensor, + tgt_loc: torch.Tensor, + src_loc: torch.Tensor, + block_size: int, + ) -> None: + if tgt_loc.numel() == 0: + return + flat = buf.reshape(-1) + tgt = tgt_loc.to(torch.int64) + src = src_loc.to(torch.int64) + tgt_page = torch.div(tgt, block_size, rounding_mode="floor") + src_page = torch.div(src, block_size, rounding_mode="floor") + tgt_pos = tgt % block_size + src_pos = src % block_size + block_stride = buf.stride(0) + token_stride = self.layout.swa_token_stride + scale_dim = self.layout.swa_scale_dim + + value_offsets = torch.arange( + token_stride, + dtype=torch.int64, + device=buf.device, + ) + tgt_value = ( + tgt_page[:, None] * block_stride + + tgt_pos[:, None] * token_stride + + value_offsets[None, :] + ) + src_value = ( + src_page[:, None] * block_stride + + src_pos[:, None] * token_stride + + value_offsets[None, :] + ) + value_rows = flat[src_value].clone() + flat[tgt_value] = value_rows + + scale_offsets = torch.arange( + scale_dim, + dtype=torch.int64, + device=buf.device, + ) + scale_base = block_size * token_stride + tgt_scale = ( + tgt_page[:, None] * block_stride + + scale_base + + tgt_pos[:, None] * scale_dim + + scale_offsets[None, :] + ) + src_scale = ( + src_page[:, None] * block_stride + + scale_base + + src_pos[:, None] * scale_dim + + scale_offsets[None, :] + ) + scale_rows = flat[src_scale].clone() + flat[tgt_scale] = scale_rows + + def _move_rows( + self, + buf: torch.Tensor, + row_bytes: int, + tgt_loc: torch.Tensor, + src_loc: torch.Tensor, + block_size: int, + ) -> None: + rows = buf.view(-1, block_size, row_bytes).reshape(-1, row_bytes) + rows[tgt_loc.long()] = rows[src_loc.long()] + + def _compressed_locs_from_token_locs( + self, + loc: torch.Tensor, + *, + ratio: int, + block_size: int, + ) -> torch.Tensor: + page = torch.div(loc.to(torch.int64), self.page_size, rounding_mode="floor") + pos = loc.to(torch.int64) % self.page_size + return page * block_size + torch.div(pos, ratio, rounding_mode="floor") + + def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor) -> None: + if tgt_loc.numel() == 0: + return + for layer_id in range(self.layer_num): + self._move_fp8_ds_mla_rows( + self.swa_kv_buffer[layer_id], + tgt_loc, + src_loc, + self.swa_block_size, + ) + buf = self.compressed_kv_buffer[layer_id] + if buf is not None: + ratio = self.layout.layer_ratio[layer_id] + block_size = self.get_compressed_block_size(layer_id) + self._move_fp8_ds_mla_rows( + buf, + self._compressed_locs_from_token_locs( + tgt_loc, ratio=ratio, block_size=block_size + ), + self._compressed_locs_from_token_locs( + src_loc, ratio=ratio, block_size=block_size + ), + block_size, + ) + for buffers, row_bytes in ( + (self.indexer_kv_buffer, self.layout.indexer_row_bytes), + ): + buf = buffers[layer_id] + if buf is not None: + ratio = self.layout.layer_ratio[layer_id] + block_size = self.get_indexer_block_size(layer_id) + self._move_rows( + buf, + row_bytes, + self._compressed_locs_from_token_locs( + tgt_loc, ratio=ratio, block_size=block_size + ), + self._compressed_locs_from_token_locs( + src_loc, ratio=ratio, block_size=block_size + ), + block_size, + ) + for buffers in (self.compressor_state_buffer, self.indexer_state_buffer): + buf = buffers[layer_id] + if buf is not None: + rows = buf.view(-1, buf.shape[-1]) + rows[tgt_loc.long()] = rows[src_loc.long()] + + def _all_buffers(self) -> list[torch.Tensor]: + out: list[torch.Tensor] = [] + for layer_id in range(self.layer_num): + out.append(self.swa_kv_buffer[layer_id]) + for buffers in ( + self.compressed_kv_buffer, + self.compressor_state_buffer, + self.indexer_kv_buffer, + self.indexer_state_buffer, + ): + buf = buffers[layer_id] + if buf is not None: + out.append(buf) + return out + + def get_kv_size_bytes(self) -> int: + return int( + sum(np.prod(buf.shape) * buf.dtype.itemsize for buf in self._all_buffers()) + ) + + def get_contiguous_buf_infos(self): + buffers = self._all_buffers() + return ( + [buf.data_ptr() for buf in buffers], + [buf.nbytes for buf in buffers], + [buf[0].nbytes for buf in buffers], + ) + + def get_layerwise_buf_info_offsets(self, start_idx=0): + offsets = [] + cursor = start_idx + for layer_id in range(self.layer_num): + layer_offsets = [cursor] + cursor += 1 + for buffers in ( + self.compressed_kv_buffer, + self.compressor_state_buffer, + self.indexer_kv_buffer, + self.indexer_state_buffer, + ): + if buffers[layer_id] is not None: + layer_offsets.append(cursor) + cursor += 1 + offsets.append(layer_offsets) + return offsets + + def get_cpu_copy(self, token_indices: list[int]) -> list[torch.Tensor]: + del token_indices + raise NotImplementedError( + "DeepSeek V4 KV cache offload is not implemented; the compressed-MQA " + "and indexer buffers are page-shaped and require page-aware indexing." + ) + + def load_cpu_copy(self, kv_cache_cpu, token_indices: list[int]) -> None: + del kv_cache_cpu, token_indices + raise NotImplementedError( + "DeepSeek V4 KV cache reload is not implemented; the compressed-MQA " + "and indexer buffers are page-shaped and require page-aware indexing." + ) diff --git a/python/tokenspeed/runtime/layers/attention/kv_cache/dsa.py b/python/tokenspeed/runtime/layers/attention/kv_cache/dsa.py new file mode 100644 index 0000000..d6db7a8 --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/kv_cache/dsa.py @@ -0,0 +1,242 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import torch +from tokenspeed_kernel.ops.quantization import quantize_fp8_with_scale + +from tokenspeed.runtime.layers.attention.configs.dsa import dsa_index_k_row_bytes +from tokenspeed.runtime.layers.attention.kv_cache.mla import ( + MLATokenToKVPool, + _get_tensor_size_bytes, +) + +_INDEX_K_FP8_GROUP_SIZE = 128 +_INDEX_K_SCALE_BYTES = torch._utils._element_size(torch.float32) + + +class DSATokenToKVPool(MLATokenToKVPool): + def __init__( + self, + *args, + index_head_dim: int, + **kwargs, + ): + self.index_head_dim = int(index_head_dim) + self.index_k_row_bytes = dsa_index_k_row_bytes(self.index_head_dim) + super().__init__(*args, **kwargs) + + with self.memory_saver_adapter.region(): + self.index_k_buffer = [ + torch.zeros( + (self.size + self.page_size, self.index_k_row_bytes), + dtype=torch.uint8, + device=self.device, + ) + for _ in range(self.layer_num) + ] + + self.index_k_data_ptrs = torch.tensor( + [buf.data_ptr() for buf in self.index_k_buffer], + dtype=torch.uint64, + device=self.device, + ) + + def _get_page_size_bytes(self): + index_size_bytes = self.index_k_row_bytes + return ( + super()._get_page_size_bytes() + + self.page_size * self.layer_num * index_size_bytes + ) + + def get_kv_size_bytes(self): + return super().get_kv_size_bytes() + _get_tensor_size_bytes(self.index_k_buffer) + + def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): + super().move_kv_cache(tgt_loc, src_loc) + if tgt_loc.numel() == 0: + return + tgt_loc_flat = tgt_loc.view(-1).long() + src_loc_flat = src_loc.view(-1).long() + for buf in self.index_k_buffer: + # Packed FP8 index-K is block-split per page, so a single token's + # bytes are NOT a contiguous row; move the FP8 values and FP32 + # scales through their block-split views instead. + fp8_view, scale_view = self._index_k_block_views(buf) + ps = self.page_size + tgt_page = tgt_loc_flat // ps + tgt_slot = tgt_loc_flat % ps + src_page = src_loc_flat // ps + src_slot = src_loc_flat % ps + fp8_view[tgt_page, tgt_slot] = fp8_view[src_page, src_slot] + scale_view[tgt_page, tgt_slot] = scale_view[src_page, src_slot] + + def has_index_k_buffer(self) -> bool: + return True + + def get_index_k_buffer(self, layer_id: int) -> torch.Tensor: + if self.layer_transfer_counter is not None: + self.layer_transfer_counter.wait_until(layer_id) + return self.index_k_buffer[layer_id] + + def set_index_k_buffer( + self, + layer_id: int, + loc: torch.Tensor, + index_k: torch.Tensor, + ) -> None: + if index_k.dtype != self.model_dtype: + index_k = index_k.to(self.model_dtype) + index_k = index_k.view(-1, self.index_head_dim) + self._set_index_k_buffer(layer_id, loc, index_k) + + def _index_k_block_views( + self, buf: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return per-page block-split views into a packed FP8 index-K buffer. + + DeepGEMM's ``fp8_paged_mqa_logits`` expects each page of ``page_size`` + tokens to be laid out as ``[page_size * head_dim FP8 values]`` followed + by ``[page_size * num_groups FP32 scales]`` (block-split), NOT a + per-token ``[fp8 | scale]`` interleave. The two ``as_strided`` views + below alias the same storage as ``buf`` so writes land in place. + + Args: + buf: Packed FP8 index-K buffer of shape ``[num_slots, row_bytes]`` + and dtype ``uint8``, where ``row_bytes == head_dim + + scale_bytes`` and ``num_slots`` is a multiple of + ``page_size``. + + Returns: + ``(fp8_view, scale_view)`` where ``fp8_view`` has shape + ``[num_pages, page_size, head_dim]`` (FP8 e4m3) and ``scale_view`` + has shape ``[num_pages, page_size, num_groups]`` (float32), both + indexed as ``view[page, slot_in_page]``. + """ + ps = self.page_size + hd = self.index_head_dim + ng = hd // _INDEX_K_FP8_GROUP_SIZE + row = hd + ng * _INDEX_K_SCALE_BYTES + num_pages = buf.shape[0] // ps + page_bytes = ps * row + flat = buf.reshape(-1) + fp8_view = torch.as_strided( + flat.view(torch.float8_e4m3fn), + (num_pages, ps, hd), + (page_bytes, hd, 1), + ) + scale_view = torch.as_strided( + flat.view(torch.float32), + (num_pages, ps, ng), + (page_bytes // 4, ng, 1), + (ps * hd) // 4, + ) + return fp8_view, scale_view + + def gather_index_k( + self, layer_id: int, slots: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Gather per-token FP8 index-K values and scales from the cache. + + The packed FP8 index-K buffer is stored block-split per page (see + :meth:`_index_k_block_views`), so the non-paged prefill + scoring kernel (``fp8_mqa_logits``), which consumes contiguous + ``(k_fp8, k_scale)`` tensors, must gather token rows through the + block-split views rather than indexing raw rows. + + Args: + layer_id: Layer whose index-K cache to read. + slots: 1D int tensor of global token slot indices to gather. + + Returns: + ``(k_fp8, k_scale)`` where ``k_fp8`` has shape + ``[num_slots, head_dim]`` (FP8 e4m3) and ``k_scale`` has shape + ``[num_slots, num_groups]`` (float32). + """ + buf = self.get_index_k_buffer(layer_id) + fp8_view, scale_view = self._index_k_block_views(buf) + slots = slots.to(torch.long) + page = slots // self.page_size + slot_in_page = slots % self.page_size + k_fp8 = fp8_view[page, slot_in_page] + k_scale = scale_view[page, slot_in_page] + return k_fp8, k_scale + + def _set_index_k_buffer( + self, + layer_id: int, + loc: torch.Tensor, + index_k: torch.Tensor, + ) -> None: + buf = self.index_k_buffer[layer_id] + index_k_fp8, index_k_scale = quantize_fp8_with_scale( + index_k, + granularity="token_group", + group_size=_INDEX_K_FP8_GROUP_SIZE, + scale_encoding="float32", + ) + + fp8_view, scale_view = self._index_k_block_views(buf) + loc = loc.to(torch.long) + page = loc // self.page_size + slot_in_page = loc % self.page_size + fp8_view[page, slot_in_page] = index_k_fp8.view(-1, self.index_head_dim) + scale_view[page, slot_in_page] = index_k_scale.view( + -1, self.index_head_dim // _INDEX_K_FP8_GROUP_SIZE + ) + + def get_contiguous_buf_infos(self): + data_ptrs, data_lens, item_lens = super().get_contiguous_buf_infos() + data_ptrs = list(data_ptrs) + data_lens = list(data_lens) + item_lens = list(item_lens) + for buf in self.index_k_buffer: + data_ptrs.append(buf.data_ptr()) + data_lens.append(buf.nbytes) + item_lens.append(buf[0].nbytes * self.page_size) + return data_ptrs, data_lens, item_lens + + def get_layerwise_buf_info_offsets(self, start_idx=0): + offsets = super().get_layerwise_buf_info_offsets(start_idx) + if self.quant_method == "per_token_head": + base_count = 3 * self.layer_num + else: + base_count = self.layer_num + return [ + layer_offsets + [start_idx + base_count + layer_id] + for layer_id, layer_offsets in enumerate(offsets) + ] + + def get_cpu_copy(self, token_indices: list[int]) -> torch.Tensor: + del token_indices + raise NotImplementedError( + "DSA KV cache offload is not implemented; sparse/indexer cache " + "buffers require page-aware layout handling." + ) + + def load_cpu_copy( + self, kv_cache_cpu: torch.Tensor, token_indices: list[int] + ) -> None: + del kv_cache_cpu, token_indices + raise NotImplementedError( + "DSA KV cache reload is not implemented; sparse/indexer cache " + "buffers require page-aware layout handling." + ) diff --git a/python/tokenspeed/runtime/layers/attention/kv_cache/flat_state_slabs.py b/python/tokenspeed/runtime/layers/attention/kv_cache/flat_state_slabs.py new file mode 100644 index 0000000..7564cb5 --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/kv_cache/flat_state_slabs.py @@ -0,0 +1,225 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""GDN/mamba2 state-slab management for the flat KV pool. + +Extracted verbatim from ``MHATokenToKVPool`` so the "MHA" (multi-head +attention KV) class no longer mixes recurrent-state bookkeeping into its +attention-KV responsibility. Behavior is byte-identical: same equalization +pre-check, same flat-GDN gate, same slab shapes/dtypes/count, same null-page +(row 0) convention, and the same ``get_state_buffers`` return values and +error cases. + +The state slabs live under the SAME page-id space as the KV pages (a single +block-id space): one ``(conv, ssm)`` pair per state LAYER, row-indexed by +page id, row 0 the never-written null page. +""" + +from __future__ import annotations + +import torch + +from tokenspeed.runtime.configs import paged_cache_spec +from tokenspeed.runtime.configs.flat_memory_plan import ( + equalized_block_size, + state_const_bytes, +) +from tokenspeed.runtime.configs.paged_cache_spec import STATE_LAYER_TYPES +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +class FlatStateSlabs: + """Owns the GDN/mamba2 conv+ssm state slabs for the flat KV pool. + + Args: + layer_types: Per-layer type labels; state layers carry the + ``STATE_LAYER_TYPES`` labels. Drives the slab pairing. + conv_state_shape / temporal_state_shape: Per-state-layer mamba2 + state tensor shapes (configs' mamba2_cache_params); ``None`` on + pure-attention models. + conv_dtype / ssm_dtype: State dtypes; default to ``default_dtype``. + default_dtype: Pool store dtype used when a state dtype is ``None``. + page_size: The (already-equalized) page size P. + size: Total token slots (the flat pool size; must be whole pages + when slabs are active). + kv_bytes_per_slot: Bytes one KV history slot occupies + (``2 * head_num * head_dim * store_dtype.itemsize``); used only + by the equalization pre-check. + + The equalization pre-check runs in ``__init__`` (same trigger, same + ``ValueError``). Device tensors are allocated lazily by ``allocate`` so + the pool can keep them inside its memory-saver "kv_cache" region. + """ + + def __init__( + self, + *, + layer_types: tuple[str, ...], + conv_state_shape: tuple[int, ...] | None, + temporal_state_shape: tuple[int, ...] | None, + conv_dtype: torch.dtype | None, + ssm_dtype: torch.dtype | None, + default_dtype: torch.dtype, + page_size: int, + size: int, + kv_bytes_per_slot: int, + ): + self._layer_types = tuple(layer_types or ()) + # Per-state-layer mamba2 shapes (configs' mamba2_cache_params); + # None on pure-attention models. + self._conv_state_shape = ( + tuple(conv_state_shape) if conv_state_shape is not None else None + ) + self._temporal_state_shape = ( + tuple(temporal_state_shape) if temporal_state_shape is not None else None + ) + self._conv_dtype = conv_dtype if conv_dtype is not None else default_dtype + self._ssm_dtype = ssm_dtype if ssm_dtype is not None else default_dtype + self.page_size = page_size + self.size = size + + # layer_id -> state pair index (the n-th state layer binds pair n). + # Derives purely from layer_types; shared by the KV skip set on the + # pool and the state-slab block below. + self._layer_state_pair: dict[int, int] = { + layer_id: pair + for pair, layer_id in enumerate( + layer_id + for layer_id, label in enumerate(self._layer_types) + if label in STATE_LAYER_TYPES + ) + } + + if ( + self._conv_state_shape is not None + and self._temporal_state_shape is not None + ): + # The flat plan packs [conv|ssm] state rows and KV rows into one + # page-id space, so P must already cover the widest constant row + # (the equalizer would otherwise inflate it behind the + # allocator's back). + equalized = equalized_block_size( + layer_types=list(self._layer_types), + kv_bytes_per_slot=kv_bytes_per_slot, + state_const_bytes=state_const_bytes( + self._conv_state_shape, + self._conv_dtype, + self._temporal_state_shape, + self._ssm_dtype, + ), + block_size=self.page_size, + alignment=1, + ) + if equalized != self.page_size: + raise ValueError( + "page_size must be pre-equalized for state layers; need " + f">= {equalized} (got {self.page_size})" + ) + + # Flat GDN predicate: ONE boolean gates both skipping per-layer KV on + # state layers and allocating the state slabs -- the plan sizing + # (registry) charges exactly full-layer KV + state rows, so the two + # decisions must never diverge. + self._flat_gdn = ( + bool(self._layer_state_pair) + and self._conv_state_shape is not None + and self._temporal_state_shape is not None + and paged_cache_spec.scheduler_ext_flat_kvcache() + ) + + self.num_pages_with_null: int | None = None + self.state_slabs: list[tuple[torch.Tensor, torch.Tensor]] = [] + + @property + def is_active(self) -> bool: + """True iff the flat-GDN gate is on (state slabs will be allocated). + + When False the pool must NOT skip per-layer KV on state layers and + must not allocate slabs (pure attention, radix ext, spec decode, or + missing state shapes).""" + return self._flat_gdn + + @property + def state_layer_ids(self) -> frozenset[int]: + """Layer ids that carry NO per-layer KV (state layers under flat + GDN). Empty unless active, so non-flat profiles keep full KV.""" + return frozenset(self._layer_state_pair) if self._flat_gdn else frozenset() + + def is_state_layer(self, layer_id: int) -> bool: + """Whether ``layer_id`` is a KV-less state layer under flat GDN.""" + return self._flat_gdn and layer_id in self._layer_state_pair + + def allocate(self, device: str) -> None: + """Allocate the ``(conv, ssm)`` slab pairs on ``device`` when active. + + Idempotent-shaped: a no-op leaving ``state_slabs == []`` when the + flat-GDN gate is off. Call inside the pool's memory-saver "kv_cache" + region so the slabs share the KV discard-on-sleep policy. + """ + if not self._flat_gdn: + self.state_slabs = [] + return + # State slabs (GDN/mamba2 conv+ssm rows): one (conv, ssm) pair per + # state LAYER (n-th state layer -> pair n), row-indexed by page id + # over the SAME page-id space as the KV pages; row 0 is the null + # page, never written -- mirrors the KV buffers' +page_size + # dummy-page convention. + assert self.size % self.page_size == 0, "flat pool size must be whole pages" + self.num_pages_with_null = self.size // self.page_size + 1 + self.state_slabs = [ + ( + torch.zeros( + (self.num_pages_with_null, *self._conv_state_shape), + dtype=self._conv_dtype, + device=device, + ), + torch.zeros( + (self.num_pages_with_null, *self._temporal_state_shape), + dtype=self._ssm_dtype, + device=device, + ), + ) + for _ in range(len(self._layer_state_pair)) + ] + logger.info( + "State slabs: %d (conv, ssm) pairs x %d page rows (row 0 = null page)", + len(self._layer_state_pair), + self.num_pages_with_null, + ) + + def get_state_buffers(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: + """(conv, ssm) state slab pair for a state layer; the n-th state + layer (within-state-label occurrence order, the slab pairing order) + binds pair n. Raises ValueError for non-state layers.""" + pair = self._layer_state_pair.get(layer_id) + if pair is None: + raise ValueError( + f"layer {layer_id} is not a state layer " + f"(layer_types={self._layer_types!r})" + ) + if not self.state_slabs: + raise ValueError( + f"layer {layer_id} is a state layer but no state " + "slabs were allocated (state shapes missing or " + "radix ext)" + ) + return self.state_slabs[pair] diff --git a/python/tokenspeed/runtime/layers/attention/kv_cache/mha.py b/python/tokenspeed/runtime/layers/attention/kv_cache/mha.py new file mode 100644 index 0000000..bef15b2 --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/kv_cache/mha.py @@ -0,0 +1,535 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from collections import Counter + +import numpy as np +import torch +from tokenspeed_kernel.ops.kvcache.triton import store_kv_cache + +from tokenspeed.runtime.configs import paged_cache_spec +from tokenspeed.runtime.configs.flat_memory_plan import occurrence_index +from tokenspeed.runtime.configs.paged_cache_spec import hybrid_slab_group_size +from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool +from tokenspeed.runtime.layers.attention.kv_cache.flat_state_slabs import ( + FlatStateSlabs, +) +from tokenspeed.runtime.layers.attention.kv_cache.utils import ( + copy_all_layer_kv_cache_tiled, + move_kv_cache_native, +) +from tokenspeed.runtime.layers.paged_attention import PagedAttention +from tokenspeed.runtime.utils import debug_timing, get_colorful_logger +from tokenspeed.runtime.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter + +logger = get_colorful_logger(__name__) + + +GB = 1024 * 1024 * 1024 + + +class MHATokenToKVPool(BaseTokenToKVPool): + def __init__( + self, + size: int, + dtype: torch.dtype, + head_num: int, + head_dim: int, + layer_num: int, + device: str, + enable_memory_saver: bool, + max_batch_size: int, + max_context_len: int, + page_size: int, + rank: int, + layer_types: tuple[str, ...] = (), + sliding_window_tokens: int | tuple[int | None, ...] | None = None, + max_scheduled_tokens: int = 0, + pd_disaggregation_enabled: bool = False, + enable_kv_cache_copy: bool = False, + enable_alt_stream: bool = True, + conv_state_shape: tuple[int, ...] | None = None, + temporal_state_shape: tuple[int, ...] | None = None, + conv_dtype: torch.dtype | None = None, + ssm_dtype: torch.dtype | None = None, + ): + super().__init__( + size, dtype, device, max_batch_size, max_context_len, page_size, rank + ) + + self.memory_saver_adapter = TorchMemorySaverAdapter.create( + enable=enable_memory_saver + ) + + self.head_num = head_num + self.head_dim = head_dim + self.layer_num = layer_num + self._layer_types = tuple(layer_types or ()) + self._pd_disaggregation_enabled = pd_disaggregation_enabled + self._slab_group_size = hybrid_slab_group_size( + self._layer_types, + sliding_window_tokens=sliding_window_tokens, + ) + # GDN/mamba2 recurrent state slabs live under this same pool object + # (one page-id space with the KV pages), but their bookkeeping is + # owned by FlatStateSlabs. Constructing it here runs the + # equalization pre-check (same trigger, same ValueError) before any + # buffer allocation; slabs themselves are allocated in + # _create_buffers inside the memory-saver region. + self._state = FlatStateSlabs( + layer_types=self._layer_types, + conv_state_shape=conv_state_shape, + temporal_state_shape=temporal_state_shape, + conv_dtype=conv_dtype, + ssm_dtype=ssm_dtype, + default_dtype=dtype, + page_size=self.page_size, + size=self.size, + kv_bytes_per_slot=2 * head_num * head_dim * self.store_dtype.itemsize, + ) + self._create_buffers() + + self.device_module = torch.get_device_module(self.device) + self.alt_stream = ( + self.device_module.Stream() + if torch.cuda.is_available() and enable_alt_stream + else None + ) + + if enable_kv_cache_copy: + self._init_kv_copy_and_warmup() + else: + self._kv_copy_config = None + + k_size, v_size = self.get_kv_size_bytes() + logger.info( + "KV Cache is allocated. K size: %.2f GB, V size: %.2f GB.", + k_size / GB, + v_size / GB, + ) + + # Publication rule lives in paged_cache_spec.publish_paged_cache_groups + # (module-attr call so tests can patch the flat-ext probe at call time). + published = paged_cache_spec.publish_paged_cache_groups( + layer_types=self._layer_types, + sliding_window_tokens=sliding_window_tokens, + page_size=page_size, + max_live_requests=max_batch_size, + max_scheduled_tokens=max_scheduled_tokens, + max_total_tokens=size, + max_context_len=max_context_len, + ) + if published is None: + self.paged_cache_group_specs = () + self.paged_cache_group_page_counts = {} + else: + specs, counts = published + self.paged_cache_group_specs = tuple(specs) + self.paged_cache_group_page_counts = counts + # Slab aliasing is only safe under the single-BlockPool ownership the + # published groups configure. + assert self._slab_group_size is None or self.paged_cache_group_specs + + def _slab_pair_index(self) -> list[int]: + """Map layer_id -> slab index: the i-th layer of every group binds + slab i (first-appearance order, as in group_specs_from_layer_types). + """ + assert self._slab_group_size is not None + assert len(self._layer_types) == self.layer_num, ( + f"hybrid slab layout: layer_types has {len(self._layer_types)} " + f"entries but layer_num={self.layer_num}" + ) + counts = Counter(self._layer_types) + assert all( + count == self._slab_group_size for count in counts.values() + ), f"hybrid slab layout: uneven groups {dict(counts)!r}" + return occurrence_index(self._layer_types) + + def _check_slab_guards(self): + """Refuse features whose per-layer buffer assumptions break when + paired layers alias the same slab tensor.""" + # kvstore is allowed (spec §6 revision): the flat L2 tier mirrors + # whole slabs byte-blind, so per-slab copies are group-safe. + if self._pd_disaggregation_enabled: + raise RuntimeError( + "hybrid slab KV layout is incompatible with PD " + "disaggregation: KV transfer registers per-layer buffer " + "pointers (get_contiguous_buf_infos), and paired layers " + "alias the same slab, so per-layer transfers would send " + "the same bytes twice and clobber the peer's pairing. Set " + "disaggregation_mode='null' or use a radix-built " + "tokenspeed_scheduler extension, which keeps the legacy " + "per-layer layout." + ) + + def _create_buffers(self): + # Tag as "kv_cache", no CPU backup: KV is discarded on sleep and rebuilt + # after wake (paging overwrites; clear_kv_buffers zeros the remapped pages). + with self.memory_saver_adapter.region(tag="kv_cache", enable_cpu_backup=False): + # Page 0 is the zero-initialized dummy page: padded tokens write + # there, and kernels may read it past valid seq_len, so its slots + # must stay finite to keep softmax well-defined. + def _alloc(): + return torch.zeros( + (self.size + self.page_size, self.head_num, self.head_dim), + dtype=self.store_dtype, + device=self.device, + ) + + # State-layer bookkeeping lives in FlatStateSlabs. The KV skip + # set below (which layers carry None KV) and the state-slab + # allocation are gated by the SAME flat-GDN predicate -- the plan + # sizing (registry) charges exactly full-layer KV + state rows, + # so the two decisions must never diverge. state_layer_ids is + # empty unless the gate is on, so non-flat profiles keep full KV. + flat_state_layers = set(self._state.state_layer_ids) + if self._state.is_active: + # Gates event_loop's retraction offload: state layers carry no + # per-layer KV, so the radix offload executor (and its host + # pool, sized for ALL layers) cannot represent this pool. + self.supports_hierarchical_kv_cache = False + + if self._slab_group_size is not None: + # Paired layers alias the same slab tensor; live rows never + # overlap (page-ownership contract in hybrid_slab_group_size). + self._check_slab_guards() + pair_index = self._slab_pair_index() + k_slabs = [_alloc() for _ in range(self._slab_group_size)] + v_slabs = [_alloc() for _ in range(self._slab_group_size)] + self.k_buffer = [ + k_slabs[pair_index[layer_id]] for layer_id in range(self.layer_num) + ] + self.v_buffer = [ + v_slabs[pair_index[layer_id]] for layer_id in range(self.layer_num) + ] + # Gates event_loop's retraction offload (built even with the + # kvstore off): per-layer host copies would alias shared slabs. + self.supports_hierarchical_kv_cache = False + logger.info( + "KV layout: hybrid slab (%d slabs x %d rows; paired " + "layers share storage; M12)", + self._slab_group_size, + self.size + self.page_size, + ) + else: + # The hybrid-slab branch above never sees state labels + # (hybrid_slab_group_size excludes them), so the skip set + # only applies here. + self.k_buffer = [ + None if layer_id in flat_state_layers else _alloc() + for layer_id in range(self.layer_num) + ] + self.v_buffer = [ + None if layer_id in flat_state_layers else _alloc() + for layer_id in range(self.layer_num) + ] + if flat_state_layers: + logger.info( + "KV layout: per-layer (%d of %d layers carry KV " + "buffers; state layers carry none)", + self.layer_num - len(flat_state_layers), + self.layer_num, + ) + else: + logger.info( + "KV layout: per-layer (%d buffers; hybrid slab " + "inactive: predicate returned None -- radix ext " + "or non-uniform/single-group layer_types)", + self.layer_num, + ) + # Pointer/stride tables carry the REAL tensors only: _kv_copy + # launches one block per data_ptrs entry (grid = numel), so a + # placeholder entry for a skipped state layer would be + # dereferenced. + real_k = [x for x in self.k_buffer if x is not None] + real_v = [x for x in self.v_buffer if x is not None] + self.k_data_ptrs = torch.tensor( + [x.data_ptr() for x in real_k], + dtype=torch.uint64, + device=self.device, + ) + self.v_data_ptrs = torch.tensor( + [x.data_ptr() for x in real_v], + dtype=torch.uint64, + device=self.device, + ) + self.data_ptrs = torch.cat([self.k_data_ptrs, self.v_data_ptrs], dim=0) + self.data_strides = torch.tensor( + [np.prod(x.shape[1:]) * x.dtype.itemsize for x in real_k + real_v], + device=self.device, + ) + + # State slabs (GDN/mamba2 conv+ssm rows) share this pool's + # memory-saver region so they follow the KV discard-on-sleep + # policy. FlatStateSlabs.allocate is a no-op (leaving + # state_slabs == []) unless the flat-GDN gate is on. + self._state.allocate(self.device) + + def _init_kv_copy_and_warmup(self): + _KV_COPY_STRIDE_THRESHOLD_LARGE = 8192 + _KV_COPY_STRIDE_THRESHOLD_MEDIUM = 4096 + _KV_COPY_TILE_SIZE_LARGE = 512 + _KV_COPY_TILE_SIZE_MEDIUM = 256 + _KV_COPY_TILE_SIZE_SMALL = 128 + _KV_COPY_NUM_WARPS_LARGE_TILE = 8 + _KV_COPY_NUM_WARPS_SMALL_TILE = 4 + + stride_bytes = int(self.data_strides[0].item()) + if stride_bytes >= _KV_COPY_STRIDE_THRESHOLD_LARGE: + bytes_per_tile = _KV_COPY_TILE_SIZE_LARGE + elif stride_bytes >= _KV_COPY_STRIDE_THRESHOLD_MEDIUM: + bytes_per_tile = _KV_COPY_TILE_SIZE_MEDIUM + else: + bytes_per_tile = _KV_COPY_TILE_SIZE_SMALL + + self._kv_copy_config = { + "bytes_per_tile": bytes_per_tile, + "byte_tiles": (stride_bytes + bytes_per_tile - 1) // bytes_per_tile, + "num_warps": ( + _KV_COPY_NUM_WARPS_SMALL_TILE + if bytes_per_tile <= _KV_COPY_TILE_SIZE_MEDIUM + else _KV_COPY_NUM_WARPS_LARGE_TILE + ), + } + + dummy_loc = torch.zeros(1, dtype=torch.int32, device=self.device) + grid = (self.data_ptrs.numel(), self._kv_copy_config["byte_tiles"]) + + copy_all_layer_kv_cache_tiled[grid]( + self.data_ptrs, + self.data_strides, + dummy_loc, + dummy_loc, + 1, + 1, + BYTES_PER_TILE=self._kv_copy_config["bytes_per_tile"], + num_warps=self._kv_copy_config["num_warps"], + num_stages=2, + ) + + def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): + # Slab layout: data_ptrs holds duplicated slab entries, so this + # broadcast re-copies rows. No callers today; re-check before wiring. + if self._kv_copy_config is None: + # Real tensors only: flat GDN state layers carry None slots. + move_kv_cache_native( + [x for x in self.k_buffer if x is not None], + [x for x in self.v_buffer if x is not None], + tgt_loc, + src_loc, + ) + else: + grid = (self.data_ptrs.numel(), self._kv_copy_config["byte_tiles"]) + copy_all_layer_kv_cache_tiled[grid]( + self.data_ptrs, + self.data_strides, + tgt_loc, + src_loc, + tgt_loc.numel(), + tgt_loc.numel(), + BYTES_PER_TILE=self._kv_copy_config["bytes_per_tile"], + num_warps=self._kv_copy_config["num_warps"], + num_stages=2, + ) + + def get_kv_size_bytes(self): + assert hasattr(self, "k_buffer") + assert hasattr(self, "v_buffer") + # Dedup by tensor identity: the slab layout aliases layers to shared + # slabs, and allocated bytes must not be double-counted. None slots + # (flat GDN state layers carry no KV) are skipped. + k_size_bytes = 0 + for k_cache in {id(t): t for t in self.k_buffer if t is not None}.values(): + k_size_bytes += np.prod(k_cache.shape) * k_cache.dtype.itemsize + v_size_bytes = 0 + for v_cache in {id(t): t for t in self.v_buffer if t is not None}.values(): + v_size_bytes += np.prod(v_cache.shape) * v_cache.dtype.itemsize + return k_size_bytes, v_size_bytes + + # for disagg + def get_contiguous_buf_infos(self): + # layer_num x [seq_len, head_num, head_dim] + # layer_num x [page_num, page_size, head_num, head_dim] + if any(x is None for x in self.k_buffer): + raise ValueError( + "flat GDN layout has no per-layer KV on state layers; " + "PD disaggregation unsupported: KV transfer registers " + "per-layer buffer pointers, and state layers carry only " + "state slabs. Set disaggregation_mode='null' or use a " + "radix-built tokenspeed_scheduler extension, which keeps " + "the full per-layer KV layout." + ) + kv_data_ptrs = [ + self._get_key_buffer(i).data_ptr() for i in range(self.layer_num) + ] + [self._get_value_buffer(i).data_ptr() for i in range(self.layer_num)] + kv_data_lens = [ + self._get_key_buffer(i).nbytes for i in range(self.layer_num) + ] + [self._get_value_buffer(i).nbytes for i in range(self.layer_num)] + kv_item_lens = [ + self._get_key_buffer(i)[0].nbytes * self.page_size + for i in range(self.layer_num) + ] + [ + self._get_value_buffer(i)[0].nbytes * self.page_size + for i in range(self.layer_num) + ] + return kv_data_ptrs, kv_data_lens, kv_item_lens + + def get_contiguous_buf_unit_lens(self): + key_units = [ + self._get_key_buffer(i)[0, 0].nbytes for i in range(self.layer_num) + ] + value_units = [ + self._get_value_buffer(i)[0, 0].nbytes for i in range(self.layer_num) + ] + return key_units + value_units + + def get_layerwise_buf_info_offsets(self, start_idx=0): + return [ + [start_idx + i * self.layer_num + layer_id for i in range(2)] + for layer_id in range(self.layer_num) + ] + + def get_cpu_copy(self, indices): + torch.cuda.synchronize() + kv_cache_cpu = [] + for layer_id in range(self.layer_num): + kv_cache_cpu.append([]) + for i in range(0, len(indices), self.offload_chunk_page_num): + chunk_indices = indices[i : i + self.offload_chunk_page_num] + k_cpu = self.k_buffer[layer_id][chunk_indices].to( + "cpu", non_blocking=True + ) + v_cpu = self.v_buffer[layer_id][chunk_indices].to( + "cpu", non_blocking=True + ) + kv_cache_cpu[-1].append([k_cpu, v_cpu]) + torch.cuda.synchronize() + return kv_cache_cpu + + def load_cpu_copy(self, kv_cache_cpu, indices): + torch.cuda.synchronize() + for layer_id in range(self.layer_num): + for i in range(0, len(indices), self.offload_chunk_page_num): + chunk_indices = indices[i : i + self.offload_chunk_page_num] + k_cpu, v_cpu = ( + kv_cache_cpu[layer_id][i // self.offload_chunk_page_num][0], + kv_cache_cpu[layer_id][i // self.offload_chunk_page_num][1], + ) + assert k_cpu.shape[0] == v_cpu.shape[0] == len(chunk_indices) + k_chunk = k_cpu.to(self.k_buffer[0].device, non_blocking=True) + v_chunk = v_cpu.to(self.v_buffer[0].device, non_blocking=True) + self.k_buffer[layer_id][chunk_indices] = k_chunk + self.v_buffer[layer_id][chunk_indices] = v_chunk + torch.cuda.synchronize() + + # Todo: different memory layout + def get_flat_data(self, indices): + # prepare a large chunk of contiguous data for efficient transfer + flatten = torch.stack( + [ + torch.stack([self.k_buffer[i][indices] for i in range(self.layer_num)]), + torch.stack([self.v_buffer[i][indices] for i in range(self.layer_num)]), + ] + ) + return flatten + + @debug_timing + def transfer(self, indices, flat_data): + # transfer prepared data from host to device + flat_data = flat_data.to(device=self.device, non_blocking=False) + k_data, v_data = flat_data[0], flat_data[1] + for i in range(self.layer_num): + self.k_buffer[i][indices] = k_data[i] + self.v_buffer[i][indices] = v_data[i] + + def _get_key_buffer(self, layer_id: int): + # for internal use of referencing + buf = self.k_buffer[layer_id] + if buf is None: + raise ValueError(f"layer {layer_id} is a state layer; it has no KV buffer") + if self.store_dtype != self.dtype: + return buf.view(self.dtype) + return buf + + def get_key_buffer(self, layer_id: int): + # note: get_key_buffer is hooked with synchronization for layer-wise KV cache loading + # it is supposed to be used only by attention backend not for information purpose + # same applies to get_value_buffer and get_kv_buffer + if self.layer_transfer_counter is not None: + self.layer_transfer_counter.wait_until(layer_id) + return self._get_key_buffer(layer_id) + + def _get_value_buffer(self, layer_id: int): + # for internal use of referencing + buf = self.v_buffer[layer_id] + if buf is None: + raise ValueError(f"layer {layer_id} is a state layer; it has no KV buffer") + if self.store_dtype != self.dtype: + return buf.view(self.dtype) + return buf + + def get_value_buffer(self, layer_id: int): + if self.layer_transfer_counter is not None: + self.layer_transfer_counter.wait_until(layer_id) + return self._get_value_buffer(layer_id) + + def get_kv_buffer(self, layer_id: int): + return self.get_key_buffer(layer_id), self.get_value_buffer(layer_id) + + @property + def state_slabs(self) -> list[tuple[torch.Tensor, torch.Tensor]]: + """(conv, ssm) state slab pairs; [] when no state slabs are active. + + Forwarding property: FlatStateSlabs owns the slabs, but the flat + host mirror and hybrid-linear-attn backend probe pool.state_slabs + directly (getattr), so keep the attribute on the pool.""" + return self._state.state_slabs + + def get_state_buffers(self, layer_id: int) -> tuple[torch.Tensor, torch.Tensor]: + """(conv, ssm) state slab pair for a state layer; the n-th state + layer (within-state-label occurrence order, the slab pairing order) + binds pair n. Raises ValueError for non-state layers.""" + return self._state.get_state_buffers(layer_id) + + def set_kv_buffer( + self, + layer: PagedAttention, + loc: torch.Tensor, + cache_k: torch.Tensor, + cache_v: torch.Tensor, + k_scale: float | None = None, + v_scale: float | None = None, + ): + layer_id = layer.layer_id + if cache_k.dtype != self.dtype: + if k_scale is not None: + cache_k.div_(k_scale) + if v_scale is not None: + cache_v.div_(v_scale) + cache_k = cache_k.to(self.dtype) + cache_v = cache_v.to(self.dtype) + if self.store_dtype != self.dtype: + cache_k = cache_k.view(self.store_dtype) + cache_v = cache_v.view(self.store_dtype) + store_kv_cache( + cache_k, cache_v, self.k_buffer[layer_id], self.v_buffer[layer_id], loc + ) diff --git a/python/tokenspeed/runtime/layers/attention/kv_cache/mla.py b/python/tokenspeed/runtime/layers/attention/kv_cache/mla.py new file mode 100644 index 0000000..67fb0a7 --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/kv_cache/mla.py @@ -0,0 +1,437 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import numpy as np +import torch + +from tokenspeed.runtime.cache.utils import ( + get_mla_kv_buffer_triton, + set_mla_kv_buffer_triton, +) +from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool +from tokenspeed.runtime.layers.attention.kv_cache.utils import ( + copy_all_layer_kv_cache_tiled, +) +from tokenspeed.runtime.layers.paged_attention import PagedAttention +from tokenspeed.runtime.utils import get_colorful_logger +from tokenspeed.runtime.utils.pdl import pdl_enabled +from tokenspeed.runtime.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter + +logger = get_colorful_logger(__name__) + +GB = 1024 * 1024 * 1024 + + +def _get_tensor_size_bytes(t: torch.Tensor | list[torch.Tensor]): + if isinstance(t, list): + return sum(_get_tensor_size_bytes(x) for x in t) + return np.prod(t.shape) * t.dtype.itemsize + + +class MLATokenToKVPool(BaseTokenToKVPool): + def __init__( + self, + size: int, + model_dtype: torch.dtype, + dtype: torch.dtype, + quant_method: str, + kv_lora_rank: int, + qk_rope_head_dim: int, + layer_num: int, + device: str, + enable_memory_saver: bool, + max_batch_size: int, + max_context_len: int, + page_size: int, + rank: int, + enable_kv_cache_copy: bool = False, + enable_alt_stream: bool = True, + ): + super().__init__( + size, dtype, device, max_batch_size, max_context_len, page_size, rank + ) + self.model_dtype = model_dtype + self.quant_method = quant_method + + self.kv_lora_rank = kv_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + self.layer_num = layer_num + self.kv_cache_dim = kv_lora_rank + qk_rope_head_dim + + self.memory_saver_adapter = memory_saver_adapter = ( + TorchMemorySaverAdapter.create(enable=enable_memory_saver) + ) + self.page_size_bytes = self._get_page_size_bytes() + + with memory_saver_adapter.region(tag="kv_cache", enable_cpu_backup=False): + # The padded page 0 is used for writing dummy outputs from padded tokens. + if self.quant_method == "per_token_head": + self.kv_buffer = [ + ( + torch.zeros( + (self.size + self.page_size, 1, kv_lora_rank), + dtype=self.store_dtype, + device=device, + ), + torch.zeros( + (self.size + self.page_size, 1, 1), + dtype=torch.float32, + device=device, + ), + torch.zeros( + (self.size + self.page_size, 1, qk_rope_head_dim), + dtype=self.model_dtype, + device=device, + ), + ) + for _ in range(layer_num) + ] + else: + self.kv_buffer = [ + torch.zeros( + (self.size + self.page_size, 1, self.kv_cache_dim), + dtype=self.store_dtype, + device=device, + ) + for _ in range(layer_num) + ] + + # Calculate data pointers and strides for all buffers + all_buffers = [] + if self.quant_method == "per_token_head": + # kv_buffer is a list of tuples (k_lora_cache, k_scale_cache, k_rope_cache) + for layer_buffers in self.kv_buffer: + # Each layer has 3 tensors + all_buffers.extend(layer_buffers) + else: + # kv_buffer is a list of single tensors + all_buffers = self.kv_buffer + + self.data_ptrs = torch.tensor( + [buf.data_ptr() for buf in all_buffers], + dtype=torch.uint64, + device=self.device, + ) + self.data_strides = torch.tensor( + [np.prod(buf.shape[1:]) * buf.dtype.itemsize for buf in all_buffers], + device=self.device, + ) + + self.device_module = torch.get_device_module(self.device) + self.alt_stream = ( + self.device_module.Stream() + if torch.cuda.is_available() and enable_alt_stream + else None + ) + + if enable_kv_cache_copy: + self._init_kv_copy_and_warmup() + else: + self._kv_copy_config = None + + def _get_page_size_bytes(self): + if self.quant_method == "per_token_head": + dim_size_bytes = ( + self.kv_lora_rank * torch._utils._element_size(self.dtype) + + self.qk_rope_head_dim * torch._utils._element_size(self.model_dtype) + + 1 * torch._utils._element_size(torch.float32) + ) + else: + dim_size_bytes = ( + self.kv_lora_rank + self.qk_rope_head_dim + ) * torch._utils._element_size(self.dtype) + return self.page_size * self.layer_num * dim_size_bytes + + def _init_kv_copy_and_warmup(self): + # Heuristics for KV copy tiling + _KV_COPY_STRIDE_THRESHOLD_LARGE = 8192 + _KV_COPY_STRIDE_THRESHOLD_MEDIUM = 4096 + _KV_COPY_TILE_SIZE_LARGE = 512 + _KV_COPY_TILE_SIZE_MEDIUM = 256 + _KV_COPY_TILE_SIZE_SMALL = 128 + _KV_COPY_NUM_WARPS_LARGE_TILE = 8 + _KV_COPY_NUM_WARPS_SMALL_TILE = 4 + + stride_bytes = int(self.data_strides[0].item()) + if stride_bytes >= _KV_COPY_STRIDE_THRESHOLD_LARGE: + bytes_per_tile = _KV_COPY_TILE_SIZE_LARGE + elif stride_bytes >= _KV_COPY_STRIDE_THRESHOLD_MEDIUM: + bytes_per_tile = _KV_COPY_TILE_SIZE_MEDIUM + else: + bytes_per_tile = _KV_COPY_TILE_SIZE_SMALL + + self._kv_copy_config = { + "bytes_per_tile": bytes_per_tile, + "byte_tiles": (stride_bytes + bytes_per_tile - 1) // bytes_per_tile, + "num_warps": ( + _KV_COPY_NUM_WARPS_SMALL_TILE + if bytes_per_tile <= _KV_COPY_TILE_SIZE_MEDIUM + else _KV_COPY_NUM_WARPS_LARGE_TILE + ), + } + + dummy_loc = torch.zeros(1, dtype=torch.int32, device=self.device) + grid = (self.data_ptrs.numel(), self._kv_copy_config["byte_tiles"]) + + copy_all_layer_kv_cache_tiled[grid]( + self.data_ptrs, + self.data_strides, + dummy_loc, + dummy_loc, + 1, + 1, + BYTES_PER_TILE=self._kv_copy_config["bytes_per_tile"], + num_warps=self._kv_copy_config["num_warps"], + num_stages=2, + ) + + def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): + if self._kv_copy_config is None: + # Native implementation for MLA + if tgt_loc.numel() == 0: + return + + tgt_loc_flat = tgt_loc.view(-1).long() + src_loc_flat = src_loc.view(-1).long() + + if self.quant_method == "per_token_head": + # kv_buffer is a list of tuples + for layer_buffers in self.kv_buffer: + # Each layer has 3 tensors: k_lora_cache, k_scale_cache, k_rope_cache + for buf in layer_buffers: + buf[tgt_loc_flat] = buf[src_loc_flat] + else: + # kv_buffer is a list of single tensors + for buf in self.kv_buffer: + buf[tgt_loc_flat] = buf[src_loc_flat] + else: + grid = (self.data_ptrs.numel(), self._kv_copy_config["byte_tiles"]) + copy_all_layer_kv_cache_tiled[grid]( + self.data_ptrs, + self.data_strides, + tgt_loc, + src_loc, + tgt_loc.numel(), + tgt_loc.numel(), + BYTES_PER_TILE=self._kv_copy_config["bytes_per_tile"], + num_warps=self._kv_copy_config["num_warps"], + num_stages=2, + ) + + def get_kv_size_bytes(self): + assert hasattr(self, "kv_buffer") + kv_size_bytes = 0 + for kv_cache in self.kv_buffer: + kv_size_bytes += _get_tensor_size_bytes(kv_cache) + return kv_size_bytes + + # for disagg + def get_contiguous_buf_infos(self): + if self.quant_method == "per_token_head": + kv_data_ptrs = [ + sub_tuple[i].data_ptr() + for i in range(3) + for sub_tuple in self.kv_buffer + ] + kv_data_lens = [ + sub_tuple[i].nbytes for i in range(3) for sub_tuple in self.kv_buffer + ] + kv_item_lens = [ + sub_tuple[i][0].nbytes * self.page_size + for i in range(3) + for sub_tuple in self.kv_buffer + ] + else: + # MLA has only one kv_buffer, so only the information of this buffer needs to be returned. + kv_data_ptrs = [self.kv_buffer[i].data_ptr() for i in range(self.layer_num)] + kv_data_lens = [self.kv_buffer[i].nbytes for i in range(self.layer_num)] + kv_item_lens = [ + self.kv_buffer[i][0].nbytes * self.page_size + for i in range(self.layer_num) + ] + return kv_data_ptrs, kv_data_lens, kv_item_lens + + def get_layerwise_buf_info_offsets(self, start_idx=0): + if self.quant_method == "per_token_head": + return [ + [start_idx + i * self.layer_num + layer_id for i in range(3)] + for layer_id in range(self.layer_num) + ] + else: + return [[start_idx + layer_id] for layer_id in range(self.layer_num)] + + def get_key_buffer(self, layer_id: int): + if self.layer_transfer_counter is not None: + self.layer_transfer_counter.wait_until(layer_id) + if self.quant_method == "per_token_head": + return self.kv_buffer[layer_id] + elif self.store_dtype != self.dtype: + return self.kv_buffer[layer_id].view(self.dtype) + else: + return self.kv_buffer[layer_id] + + def get_value_buffer(self, layer_id: int): + if self.layer_transfer_counter is not None: + self.layer_transfer_counter.wait_until(layer_id) + if self.quant_method == "per_token_head": + return self.kv_buffer[layer_id][:2] + elif self.store_dtype != self.dtype: + return self.kv_buffer[layer_id][..., : self.kv_lora_rank].view(self.dtype) + else: + return self.kv_buffer[layer_id][..., : self.kv_lora_rank] + + def get_kv_buffer(self, layer_id: int): + return self.get_key_buffer(layer_id), self.get_value_buffer(layer_id) + + def set_kv_buffer( + self, + layer: PagedAttention, + loc: torch.Tensor, + cache_k: torch.Tensor, + cache_v: torch.Tensor, + k_scale: float | None = None, + v_scale: float | None = None, + ): + layer_id = layer.layer_id + if self.quant_method == "per_token_head": + k_lora = cache_k[..., : self.kv_lora_rank].float() + k_rope = cache_k[..., self.kv_lora_rank :].float() + scale = k_lora.abs().amax(dim=-1, keepdim=True).clamp(1e-26) / 448.0 + k_lora = (k_lora / scale).to(torch.float8_e4m3fn) + k_rope = (k_rope / scale).to(self.model_dtype) + self.kv_buffer[layer_id][0][loc] = k_lora.view(self.store_dtype) + self.kv_buffer[layer_id][1][loc] = scale + self.kv_buffer[layer_id][2][loc] = k_rope + else: + self.kv_buffer[layer_id][loc] = cache_k + + def set_mla_kv_buffer( + self, + layer: PagedAttention, + loc: torch.Tensor, + cache_k_nope: torch.Tensor, + cache_k_rope: torch.Tensor, + ): + layer_id = layer.layer_id + if self.quant_method == "per_token_head": + k_lora = cache_k_nope.float() + k_rope = cache_k_rope.float() + scale = k_lora.abs().amax(dim=-1, keepdim=True).clamp(1e-26) / 448.0 + k_lora = (k_lora / scale).to(torch.float8_e4m3fn) + k_rope = (k_rope / scale).to(self.model_dtype) + self.kv_buffer[layer_id][0][loc] = k_lora.view(self.store_dtype) + self.kv_buffer[layer_id][1][loc] = scale + self.kv_buffer[layer_id][2][loc] = k_rope + else: + if cache_k_nope.dtype != self.dtype: + cache_k_nope = cache_k_nope.to(self.dtype) + cache_k_rope = cache_k_rope.to(self.dtype) + if self.store_dtype != self.dtype: + cache_k_nope = cache_k_nope.view(self.store_dtype) + cache_k_rope = cache_k_rope.view(self.store_dtype) + + set_mla_kv_buffer_triton( + self.kv_buffer[layer_id], + loc, + cache_k_nope, + cache_k_rope, + enable_pdl=pdl_enabled(), + ) + + def get_mla_kv_buffer( + self, + layer: PagedAttention, + loc: torch.Tensor, + dst_dtype: torch.dtype | None = None, + ): + layer_id = layer.layer_id + dst_dtype = dst_dtype or self.dtype + + if self.quant_method == "per_token_head": + k_lora_cache, k_scale_cache, k_rope_cache = self.kv_buffer[layer_id] + k_lora = k_lora_cache[loc].view(self.dtype).float() + k_scale = k_scale_cache[loc] + k_rope = k_rope_cache[loc].float() + cache_k_nope = (k_lora * k_scale).to(dst_dtype).contiguous() + cache_k_rope = (k_rope * k_scale).to(dst_dtype).contiguous() + return cache_k_nope, cache_k_rope + + kv_buffer = self.get_key_buffer(layer_id) + cache_k_nope = torch.empty( + (loc.shape[0], 1, self.kv_lora_rank), + dtype=dst_dtype, + device=kv_buffer.device, + ) + cache_k_rope = torch.empty( + (loc.shape[0], 1, self.qk_rope_head_dim), + dtype=dst_dtype, + device=kv_buffer.device, + ) + get_mla_kv_buffer_triton( + kv_buffer, loc, cache_k_nope, cache_k_rope, enable_pdl=pdl_enabled() + ) + return cache_k_nope, cache_k_rope + + def get_cpu_copy(self, token_indices: list[int]) -> torch.Tensor: + torch.cuda.synchronize() + kv_cache_cpu = [] + for layer_id in range(self.layer_num): + kv_cache_cpu.append([]) + for i in range(0, len(token_indices), self.offload_chunk_page_num): + chunk_indices = token_indices[i : i + self.offload_chunk_page_num] + if self.quant_method == "per_token_head": + kv_cache_cpu[-1].append( + [ + buffer[chunk_indices].to("cpu", non_blocking=True) + for buffer in self.kv_buffer[layer_id] + ] + ) + else: + kv_cpu = self.kv_buffer[layer_id][chunk_indices].to( + "cpu", non_blocking=True + ) + kv_cache_cpu[-1].append([kv_cpu]) + torch.cuda.synchronize() + return kv_cache_cpu + + def load_cpu_copy( + self, kv_cache_cpu: torch.Tensor, token_indices: list[int] + ) -> None: + torch.cuda.synchronize() + for layer_id in range(self.layer_num): + for i in range(0, len(token_indices), self.offload_chunk_page_num): + chunk_indices = token_indices[i : i + self.offload_chunk_page_num] + if self.quant_method == "per_token_head": + for j in range(3): + t = kv_cache_cpu[layer_id][i // self.offload_chunk_page_num][j] + assert t.shape[0] == len(chunk_indices) + self.kv_buffer[layer_id][j][chunk_indices] = t.to( + self.kv_buffer[0][0].device, non_blocking=True + ) + else: + kv_cpu = kv_cache_cpu[layer_id][i // self.offload_chunk_page_num][0] + assert kv_cpu.shape[0] == len( + chunk_indices + ), f"kv_cpu.shape[0] {kv_cpu.shape[0]} != len(chunk_indices) {len(chunk_indices)}" + kv_chunk = kv_cpu.to(self.kv_buffer[0].device, non_blocking=True) + self.kv_buffer[layer_id][chunk_indices] = kv_chunk + torch.cuda.synchronize() diff --git a/python/tokenspeed/runtime/layers/attention/kv_cache/utils.py b/python/tokenspeed/runtime/layers/attention/kv_cache/utils.py new file mode 100644 index 0000000..9155a23 --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/kv_cache/utils.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the FluentLLM project +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import torch +import triton +import triton.language as tl + + +def move_kv_cache_native( + k_buffer: list[torch.Tensor], + v_buffer: list[torch.Tensor], + tgt_loc: torch.Tensor, + src_loc: torch.Tensor, +): + if tgt_loc.numel() == 0: + return + + tgt_loc_flat = tgt_loc.view(-1).long() + src_loc_flat = src_loc.view(-1).long() + for k_cache, v_cache in zip(k_buffer, v_buffer): + k_cache[tgt_loc_flat] = k_cache[src_loc_flat] + v_cache[tgt_loc_flat] = v_cache[src_loc_flat] + + +@triton.jit +def copy_all_layer_kv_cache_tiled( + data_ptrs, + strides, + tgt_loc_ptr, + src_loc_ptr, + num_locs, + num_locs_upper: tl.constexpr, + BYTES_PER_TILE: tl.constexpr, +): + """2D tiled kernel. Safe for in-place copy.""" + bid = tl.program_id(0) + tid = tl.program_id(1) + + stride = tl.load(strides + bid) + base_ptr = tl.load(data_ptrs + bid) + base_ptr = tl.cast(base_ptr, tl.pointer_type(tl.uint8)) + + byte_off = tid * BYTES_PER_TILE + tl.arange(0, BYTES_PER_TILE) + mask_byte = byte_off < stride + tl.multiple_of(byte_off, 16) + + loc_idx = tl.arange(0, num_locs_upper) + mask_loc = loc_idx < num_locs + + src = tl.load(src_loc_ptr + loc_idx, mask=mask_loc, other=0) + tgt = tl.load(tgt_loc_ptr + loc_idx, mask=mask_loc, other=0) + + src_ptr = base_ptr + src[:, None] * stride + byte_off[None, :] + tgt_ptr = base_ptr + tgt[:, None] * stride + byte_off[None, :] + + mask = mask_loc[:, None] & mask_byte[None, :] + vals = tl.load(src_ptr, mask=mask) + tl.store(tgt_ptr, vals, mask=mask) diff --git a/python/tokenspeed/runtime/layers/attention/linear/causal_conv1d.py b/python/tokenspeed/runtime/layers/attention/linear/causal_conv1d.py new file mode 100755 index 0000000..f586a09 --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/linear/causal_conv1d.py @@ -0,0 +1,1108 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright (c) 2024, Tri Dao +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + + +import numpy as np +import torch +import triton +import triton.language as tl + +PAD_SLOT_ID = -1 + + +@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, + # 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: + 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() # 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: torch.Tensor | None, + conv_states: torch.Tensor, + query_start_loc: torch.Tensor, + cache_indices: torch.Tensor | None = None, + has_initial_state: torch.Tensor | None = None, + activation: str | None = "silu", + pad_slot_id: int = PAD_SLOT_ID, + metadata=None, + validate_data=False, + **kwargs, +): + """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 + args = nums_dict + batch_ptr = metadata.batch_ptr + token_chunk_offset_ptr = metadata.token_chunk_offset_ptr + else: + seq_lens_cpu = kwargs.get("seq_lens_cpu") + if seq_lens_cpu is not None: + seqlens = np.asarray(seq_lens_cpu) + 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: + # 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) + 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): + nums = -(-seqlens // META["BLOCK_M"]) # ceil-div, numpy array + tot = int(nums.sum()) + + mlist = np.repeat(np.arange(len(nums)), nums) + # offsetlist[i] = local chunk index within its sequence + offsetlist = np.arange(tot) - np.repeat(np.cumsum(nums) - nums, nums) + mlist_len = mlist.shape[0] + + 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) + + combined_np = np.stack([mlist, offsetlist]).astype(np.int32, copy=False) + combined_cpu = torch.from_numpy(combined_np).pin_memory() + META["batch_ptr"][:mlist_len].copy_(combined_cpu[0], non_blocking=True) + META["token_chunk_offset_ptr"][:mlist_len].copy_( + combined_cpu[1], non_blocking=True + ) + + 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, + 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, + output_state_indices_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, + # 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_output_state_indices_seq: tl.constexpr, + stride_output_state_indices_step: 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, + HAS_OUTPUT_STATE_INDICES: 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: + 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: + 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) + + x_base = x_ptr + (idx_seq * stride_x_seq) + (idx_feats * stride_x_dim) # [BLOCK_N] + + if not SAVE_INTERMEDIATE and not HAS_OUTPUT_STATE_INDICES: + # STEP 2: update conv_state in place. Speculative verify uses + # SAVE_INTERMEDIATE and scatters the accepted intermediate window after + # verification, so writing the real conv_state here is both wrong and + # can address beyond its width-1 storage for multi-token verify. + 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_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: [batch_position, step, dim, win(K-1)] + # Use idx_seq (batch position) instead of conv_state_batch_coord + # (pool index) so the intermediate cache can be sized to + # max_batch_size rather than the full mamba pool. + base_ptr = ( + intermediate_conv_window_ptr + + idx_seq * 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) + if HAS_OUTPUT_STATE_INDICES: + output_state_idx = tl.load( + output_state_indices_ptr + + idx_seq * stride_output_state_indices_seq + + idx_token * stride_output_state_indices_step + ).to(tl.int64) + if output_state_idx >= 0: + output_base = ( + conv_state_ptr + + output_state_idx * stride_conv_state_seq + + idx_feats * stride_conv_state_dim + ) + if KERNEL_WIDTH >= 2: + tl.store(output_base + 0 * stride_conv_state_tok, col0, mask=mask_w) + if KERNEL_WIDTH >= 3: + tl.store(output_base + 1 * stride_conv_state_tok, col1, mask=mask_w) + if KERNEL_WIDTH >= 4: + tl.store(output_base + 2 * stride_conv_state_tok, col2, mask=mask_w) + + +def causal_conv1d_update( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None = None, + activation: bool | str | None = None, + cache_seqlens: torch.Tensor | None = None, + conv_state_indices: torch.Tensor | None = None, + num_accepted_tokens: torch.Tensor | None = None, + intermediate_conv_window: torch.Tensor | None = None, + output_state_indices: torch.Tensor | None = 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 + 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 + + 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 + ) + if output_state_indices is not None or intermediate_conv_window is not None: + state_len = width - 1 + else: + 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 + if output_state_indices is not None: + stride_output_state_indices_seq, stride_output_state_indices_step = ( + output_state_indices.stride(0), + output_state_indices.stride(1), + ) + else: + stride_output_state_indices_seq = stride_output_state_indices_step = 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, + output_state_indices if output_state_indices 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_output_state_indices_seq, + stride_output_state_indices_step, + 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, + HAS_OUTPUT_STATE_INDICES=output_state_indices is not None, + ) + if unsqueeze: + out = out.squeeze(-1) + return out diff --git a/python/tokenspeed/runtime/layers/attention/linear/fused_sigmoid_gating_recurrent.py b/python/tokenspeed/runtime/layers/attention/linear/fused_sigmoid_gating_recurrent.py new file mode 100755 index 0000000..0bb9c41 --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/linear/fused_sigmoid_gating_recurrent.py @@ -0,0 +1,341 @@ +# SPDX-License-Identifier: MIT +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + + +import torch +import triton +import triton.language as tl + + +@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, + # Parameters for target_verify support (unused for decode) + intermediate_states_buffer, + cache_steps, + output_state_indices, + retrieve_parent_token_ptr, + stride_retrieve_parent_token_seq: tl.constexpr, + stride_retrieve_parent_token_token: tl.constexpr, + # ================================================ + scale, + T, + stride_q, + stride_k, + stride_v, + stride_b, + NP2_T: tl.constexpr, + 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, + # Optional flags for target_verify support (default False for decode) + DISABLE_STATE_UPDATE: tl.constexpr = False, + CACHE_INTERMEDIATE_STATES: tl.constexpr = False, + HAS_OUTPUT_STATE_INDICES: tl.constexpr = False, + HAS_EAGLE_TREE_CUSTOM_ATTN_MASK: tl.constexpr = False, +): + """ + 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 * stride_q + i_h * K + o_k + p_k = k + bos * stride_k + i_h * K + o_k + p_v = v + bos * stride_v + i_hv * V + o_v + p_b = b + bos * stride_b + 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) + + # Prepare intermediate state cache index if enabled + cache_idx = -1 + if CACHE_INTERMEDIATE_STATES: + cache_idx = i_n + + step_idx = 0 + 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) + + # Cache intermediate states if enabled + if HAS_OUTPUT_STATE_INDICES: + out_idx = tl.load(output_state_indices + i_n * T + step_idx).to(tl.int64) + if out_idx >= 0: + output_ptr = ( + h0_source + + out_idx * HV * K * V + + i_hv * K * V + + o_k[:, None] * V + + o_v[None, :] + ) + tl.store(output_ptr, b_h.to(output_ptr.dtype.element_ty), mask=mask_h) + elif CACHE_INTERMEDIATE_STATES: + if cache_idx >= 0: + 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 + + # Update pointers for next timestep + p_q += stride_q + p_k += stride_k + p_v += stride_v + p_b += stride_b + p_o += HV * V + p_a += HV + + # Store final state back to h0_source with bounds checking + if not DISABLE_STATE_UPDATE: + 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) + + +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: float | None = None, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.Tensor | None = None, + # Optional parameters for target_verify support + disable_state_update: bool = False, + intermediate_states_buffer: torch.Tensor | None = None, + cache_steps: int | None = None, + output_state_indices: torch.Tensor | None = None, + retrieve_parent_token: torch.Tensor | None = 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. + + Supports both decode and target_verify modes: + - decode: standard single-step update with state write-back + - target_verify: multi-step with intermediate state caching, optional tree attention, + and optional state update disable + """ + B, T, H, K, V = *k.shape, v.shape[-1] + stride_q = q.stride()[1] + stride_k = k.stride()[1] + stride_v = v.stride()[1] + stride_b = b.stride()[-2] + 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), 32) + 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) + + # Prepare retrieve_parent_token strides + if retrieve_parent_token is not None: + stride_retrieve_parent_token_seq = retrieve_parent_token.stride(0) + stride_retrieve_parent_token_token = retrieve_parent_token.stride(1) + else: + stride_retrieve_parent_token_seq = 0 + stride_retrieve_parent_token_token = 0 + + NP2_T = triton.next_power_of_2(T) + + 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, + intermediate_states_buffer=intermediate_states_buffer, + cache_steps=0 if cache_steps is None else cache_steps, + output_state_indices=output_state_indices, + retrieve_parent_token_ptr=retrieve_parent_token, + stride_retrieve_parent_token_seq=stride_retrieve_parent_token_seq, + stride_retrieve_parent_token_token=stride_retrieve_parent_token_token, + scale=scale, + T=T, + stride_q=stride_q, + stride_k=stride_k, + stride_v=stride_v, + stride_b=stride_b, + NP2_T=NP2_T, + B=B, + H=H, + HV=HV, + K=K, + V=V, + BK=BK, + BV=BV, + USE_INITIAL_STATE=initial_state_source is not None, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + IS_VARLEN=cu_seqlens is not None, + DISABLE_STATE_UPDATE=disable_state_update, + CACHE_INTERMEDIATE_STATES=intermediate_states_buffer is not None, + HAS_OUTPUT_STATE_INDICES=output_state_indices is not None, + HAS_EAGLE_TREE_CUSTOM_ATTN_MASK=retrieve_parent_token is not None, + num_warps=num_warps, + num_stages=num_stages, + ) + o = o.squeeze(0) + return o diff --git a/python/tokenspeed/runtime/layers/attention/linear/gdn.py b/python/tokenspeed/runtime/layers/attention/linear/gdn.py new file mode 100644 index 0000000..b511047 --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/linear/gdn.py @@ -0,0 +1,71 @@ +# SPDX-License-Identifier: MIT +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import torch +import triton +import triton.language as tl + + +@triton.jit +def fused_gdn_gating_kernel( + g, + A_log, + a, + dt_bias, + seq_len, + NUM_HEADS: tl.constexpr, + beta: tl.constexpr, + threshold: tl.constexpr, + BLK_HEADS: tl.constexpr, +): + i_b, i_s, i_d = tl.program_id(0), tl.program_id(1), tl.program_id(2) + head_off = i_d * BLK_HEADS + tl.arange(0, BLK_HEADS) + off = i_b * seq_len * NUM_HEADS + i_s * NUM_HEADS + head_off + mask = head_off < NUM_HEADS + blk_A_log = tl.load(A_log + head_off, mask=mask) + blk_a = tl.load(a + off, mask=mask) + blk_bias = tl.load(dt_bias + head_off, mask=mask) + x = blk_a.to(tl.float32) + blk_bias.to(tl.float32) + softplus_x = tl.where( + beta * x <= threshold, (1 / beta) * tl.log(1 + tl.exp(beta * x)), x + ) + blk_g = -tl.exp(blk_A_log.to(tl.float32)) * softplus_x + tl.store(g + off, blk_g.to(g.dtype.element_ty), mask=mask) + + +def fused_gdn_gating( + A_log: torch.Tensor, + a: torch.Tensor, + dt_bias: torch.Tensor, + beta: float = 1.0, + threshold: float = 20.0, +) -> torch.Tensor: + batch, num_heads = a.shape + seq_len = 1 + grid = (batch, seq_len, triton.cdiv(num_heads, 8)) + g = torch.empty_like(a, dtype=torch.float32) + fused_gdn_gating_kernel[grid]( + g, A_log, a, dt_bias, seq_len, num_heads, beta, threshold, 8, num_warps=1 + ) + return g diff --git a/python/tokenspeed/runtime/layers/attention/linear/layernorm_gated.py b/python/tokenspeed/runtime/layers/attention/linear/layernorm_gated.py new file mode 100755 index 0000000..fe0d65f --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/linear/layernorm_gated.py @@ -0,0 +1,256 @@ +# SPDX-License-Identifier: MIT +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# SPDX-FileCopyrightText: Copyright (c) 2024, Tri Dao +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +# 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 torch +import triton +import triton.language as tl + + +@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 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 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/tokenspeed/runtime/layers/attention/linear/mamba_state_scatter_triton.py b/python/tokenspeed/runtime/layers/attention/linear/mamba_state_scatter_triton.py new file mode 100644 index 0000000..756d15f --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/linear/mamba_state_scatter_triton.py @@ -0,0 +1,264 @@ +# SPDX-License-Identifier: MIT +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Fused Triton kernels for Mamba state copy and zero operations.""" + +import torch +import triton +import triton.language as tl + + +@triton.jit +def _mamba_state_snapshot_kernel( + pool_ptr, + src_indices_ptr, # [num_valid] + dst_indices_ptr, # [num_valid] + cache_lengths_ptr, # [num_valid] or nullptr (0 when page_size==0) + page_size, # 0 means no page filtering + elem_per_entry: tl.constexpr, + layer_stride, + req_stride, + pool_size, + BLOCK_SIZE: tl.constexpr, +): + """ + In-place copy kernel: pool[:, dst[i], :] = pool[:, src[i], :] + Skips copy if page_size > 0 and cache_lengths[i] % page_size != 0. + + Grid: (num_valid, num_layers) — loops over elem_per_entry internally. + Invalid entries early-return wasting only 1 block instead of + ceil(elem_per_entry / BLOCK_SIZE) blocks. + """ + pid_req = tl.program_id(0) + pid_layer = tl.program_id(1).to(tl.int64) + + src_idx = tl.load(src_indices_ptr + pid_req).to(tl.int64) + dst_idx = tl.load(dst_indices_ptr + pid_req).to(tl.int64) + + # Skip self-copy (no-op) + if src_idx == dst_idx: + return + + # Page-boundary filter: skip if not aligned + if page_size > 0: + cl = tl.load(cache_lengths_ptr + pid_req).to(tl.int64) + if cl % page_size != 0: + return + + # Bounds check + if not ( + (src_idx >= 0) & (src_idx < pool_size) & (dst_idx >= 0) & (dst_idx < pool_size) + ): + return + + src_offset = pid_layer * layer_stride + src_idx * req_stride + dst_offset = pid_layer * layer_stride + dst_idx * req_stride + + for start in tl.static_range(0, elem_per_entry, BLOCK_SIZE): + offsets = start + tl.arange(0, BLOCK_SIZE) + mask = offsets < elem_per_entry + data = tl.load(pool_ptr + src_offset + offsets, mask=mask) + tl.store(pool_ptr + dst_offset + offsets, data, mask=mask) + + +def fused_mamba_state_copy( + pool: torch.Tensor, # [num_layers, pool_size, *state_shape] or [pool_size, *state_shape] + src_indices: torch.Tensor, # [num_valid] + dst_indices: torch.Tensor, # [num_valid] + cache_lengths: torch.Tensor | None = None, # [num_valid], for page filter + page_size: int = 0, # 0 means no page filtering + single_layer: bool | None = None, +): + """ + Copy mamba states: pool[:, dst_indices[i], :] = pool[:, src_indices[i], :] + + Handles both COW copy and checkpoint snapshot. Invalid indices (< 0 or + >= pool_size) are skipped inside the kernel. When page_size > 0 and + cache_lengths is provided, also skips entries where + cache_lengths[i] % page_size != 0. + + Supports both 3D pool tensors ``[num_layers, pool_size, *state_shape]`` + and 2D per-layer slices ``[pool_size, *state_shape]`` (which may be + non-contiguous views into a larger cache). For 2D inputs the kernel + is launched with ``num_layers=1``. + Per-layer slices with rank > 2 (for example + ``[pool_size, channels, state_len]``) must pass ``single_layer=True`` so + the leading dimension is interpreted as the slot dimension, not as a layer + dimension. + + Args: + pool: State tensor, either 3D [num_layers, pool_size, *state_shape] + or 2D [pool_size, *state_shape]. + src_indices: Source slot indices [num_valid], int32 or int64. + dst_indices: Destination slot indices [num_valid], int32 or int64. + single_layer: Interpret ``pool`` as a per-layer ``[pool_size, ...]`` + slice. Defaults to True only for rank-2 tensors for compatibility. + cache_lengths: Per-entry cache lengths for page-boundary filtering. + page_size: When > 0, only copy entries where cache_lengths[i] is + aligned to page_size. Set to 0 to disable filtering (used by + COW copy where all valid entries must be copied). + """ + num_valid = src_indices.shape[0] + if num_valid == 0: + return + + if not pool.is_cuda: + raise ValueError("fused_mamba_state_copy only supports CUDA tensors.") + if pool.ndim < 2: + raise ValueError(f"pool must be at least 2D, got {pool.ndim}D") + if src_indices.shape[0] != dst_indices.shape[0]: + raise ValueError( + f"indices length mismatch: {src_indices.shape[0]} vs {dst_indices.shape[0]}" + ) + + if single_layer is None: + single_layer = pool.ndim == 2 + + if single_layer: + num_layers = 1 + pool_size = pool.shape[0] + elem_per_entry = pool.numel() // pool_size + layer_stride = 0 + req_stride = pool.stride(0) + else: + num_layers = pool.shape[0] + pool_size = pool.shape[1] + elem_per_entry = pool.numel() // (num_layers * pool_size) + layer_stride = pool.stride(0) + req_stride = pool.stride(1) + + if not src_indices.is_contiguous(): + src_indices = src_indices.contiguous() + if not dst_indices.is_contiguous(): + dst_indices = dst_indices.contiguous() + src_indices = src_indices.to(torch.int32) + dst_indices = dst_indices.to(torch.int32) + + if page_size > 0 and cache_lengths is not None: + cache_lengths = cache_lengths.to(torch.int32) + else: + cache_lengths = src_indices # unused; kernel skips when page_size==0 + page_size = 0 + + BLOCK_SIZE = 8192 + grid = (num_valid, num_layers) + + _mamba_state_snapshot_kernel[grid]( + pool, + src_indices, + dst_indices, + cache_lengths, + page_size, + elem_per_entry, + layer_stride, + req_stride, + pool_size, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=8, + ) + + +@triton.jit +def _mamba_state_zero_kernel( + pool_ptr, + indices_ptr, # [bs] — indices to zero; negative values are skipped + elem_per_entry: tl.constexpr, + layer_stride, + req_stride, + pool_size, + BLOCK_SIZE: tl.constexpr, +): + """ + Zero kernel: pool[:, indices[i], :] = 0 + Skips entries where indices[i] < 0 or indices[i] >= pool_size. + + Grid: (bs, num_layers) — loops over elem_per_entry internally. + """ + pid_req = tl.program_id(0) + pid_layer = tl.program_id(1).to(tl.int64) + + idx = tl.load(indices_ptr + pid_req).to(tl.int64) + + # Skip invalid entries (negative sentinel from torch.where) + if (idx < 0) | (idx >= pool_size): + return + + dst_offset = pid_layer * layer_stride + idx * req_stride + + zero_val = tl.zeros([BLOCK_SIZE], dtype=pool_ptr.dtype.element_ty) + for start in tl.static_range(0, elem_per_entry, BLOCK_SIZE): + offsets = start + tl.arange(0, BLOCK_SIZE) + mask = offsets < elem_per_entry + tl.store(pool_ptr + dst_offset + offsets, zero_val, mask=mask) + + +def fused_mamba_state_zero( + pool: torch.Tensor, # [num_layers, pool_size, *state_shape] + indices: torch.Tensor, # [bs] — slots to zero; negative values skipped inside kernel +): + """ + Zero mamba states: pool[:, indices[i], :] = 0 for valid indices. + + Invalid indices (< 0) are skipped inside the kernel, avoiding any + CPU-GPU synchronization from boolean indexing or .any() checks. + + Args: + pool: State tensor [num_layers, pool_size, *state_shape], must be contiguous. + indices: Slot indices [bs], int64. Negative values are treated as invalid + and skipped (no-op). + """ + bs = indices.shape[0] + if bs == 0: + return + + if not pool.is_cuda: + raise ValueError("fused_mamba_state_zero only supports CUDA tensors.") + if not pool.is_contiguous(): + raise ValueError("pool tensor must be contiguous") + if pool.ndim < 2: + raise ValueError(f"pool must be at least 2D, got {pool.ndim}D") + + num_layers = pool.shape[0] + pool_size = pool.shape[1] + elem_per_entry = pool.numel() // (num_layers * pool_size) + + layer_stride = pool.stride(0) + req_stride = pool.stride(1) + + if not indices.is_contiguous(): + indices = indices.contiguous() + + BLOCK_SIZE = 8192 + grid = (bs, num_layers) + + _mamba_state_zero_kernel[grid]( + pool, + indices, + elem_per_entry, + layer_stride, + req_stride, + pool_size, + BLOCK_SIZE=BLOCK_SIZE, + num_warps=8, + ) diff --git a/python/tokenspeed/runtime/layers/attention/mm_encoder_attention.py b/python/tokenspeed/runtime/layers/attention/mm_encoder_attention.py new file mode 100644 index 0000000..cbce7c7 --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/mm_encoder_attention.py @@ -0,0 +1,560 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Cache-less multi-headed attention used by multimodal encoders. + +Encode-only attention layer used by vision and audio towers, plus its backend +dispatch table. Backends are free functions (not +``AttentionBackend`` subclasses) because the vision encoder is single-shot +with no KV cache, no decode/extend split, and no graph capture protocol -- +the ``AttentionBackend`` ABC's prefill/decode/extend lifecycle does not +apply. This file is therefore kept out of ``backends/`` (the home of +``AttentionBackend`` subclasses) and lives at the ``layers/attention/`` top +level alongside ``registry.py`` / ``utils.py``. +""" + +from __future__ import annotations + +import functools +import logging +from collections.abc import Callable +from typing import Any + +import torch +import torch.nn as nn +from einops import rearrange +from tokenspeed_kernel.platform import current_platform + +from tokenspeed.runtime.distributed import utils as dist_utils +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.layers.linear import ( + QKVParallelLinear, + RowParallelLinear, +) +from tokenspeed.runtime.layers.quantization import QuantizationConfig +from tokenspeed.runtime.layers.rotary_embedding import apply_rotary_pos_emb_native +from tokenspeed.runtime.utils import add_prefix, round_up + +logger = logging.getLogger(__name__) + +_platform = current_platform() +_is_nvidia = _platform.is_nvidia +_is_amd = _platform.is_amd + +if _is_nvidia: + from tokenspeed_kernel.ops.attention.flash_attn import flash_attn_varlen_func + from tokenspeed_kernel.ops.attention.flashinfer import ( + cudnn_batch_prefill_with_kv_cache, + ) + +from tokenspeed_kernel.ops.attention.triton.context import context_attention_fwd +from tokenspeed_kernel.ops.attention.triton.qkv_rotary import ( + packed_qkv_complex_rotary, + packed_qkv_neox_rotary, +) + +# CUDA-graph bucketing for the cuDNN vision prefill backend: batch and max +# seqlen are quantized so a small set of captured graphs covers the request +# distribution. The consts are consumed by VLM tower models, not by +# ``MultimodalEncoderAttention`` itself. +VIT_CUDNN_WORKSPACE_BYTES = 128 * 1024 * 1024 +VIT_CUDNN_BATCH_BUCKETS: tuple[int, ...] = (8, 16, 32, 64) +VIT_CUDNN_SEQLEN_BUCKETS: tuple[int, ...] = (4096, 8192, 16384, 32768, 65536, 131072) + + +def round_up_to_bucket(value: int, buckets: tuple[int, ...]) -> int: + """Smallest bucket >= value; values past the last bucket round up to a + multiple of it. Used by vision tower code to pad batch size and max-seqlen + into a finite set of captured cuDNN graph shapes. + """ + if value <= 0: + return buckets[0] + for bucket in buckets: + if bucket >= value: + return bucket + return round_up(value, buckets[-1]) + + +# === Backend dispatch === +# The dispatcher always passes the full kwarg set (cu_seqlens / bsz / seq_len / +# softmax_scale / max_seqlen / sequence_lengths / workspace_buffer); each +# backend declares the subset it uses and absorbs the rest via ``**_``. + + +def _varlen_metadata( + cu_seqlens: torch.Tensor | None, + bsz: int, + seq_len: int, + *, + device: torch.device, + max_seqlen: int | None, +) -> tuple[torch.Tensor, torch.Tensor, int]: + """Resolve cu_seqlens / seq_lens / max_seqlen shared by the varlen backends. + + ``max_seqlen`` is honored when the caller supplies it (the capture-safe + path); only the eager fallback derives it via ``.item()``, which forces a + GPU->CPU sync that is illegal inside a captured CUDA graph. Deriving it + once here keeps every varlen backend capture-safe instead of each kernel + wrapper re-deriving (and re-syncing) it. + """ + if cu_seqlens is None: + cu_seqlens = torch.arange( + 0, (bsz + 1) * seq_len, step=seq_len, dtype=torch.int32, device=device + ) + else: + cu_seqlens = cu_seqlens.to(dtype=torch.int32, device=device) + seq_lens = cu_seqlens[1:] - cu_seqlens[:-1] + if max_seqlen is None: + max_seqlen = int(seq_lens.max().item()) + return cu_seqlens, seq_lens, int(max_seqlen) + + +def vision_attn_triton( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + cu_seqlens: torch.Tensor | None, + bsz: int, + seq_len: int, + softmax_scale: float | None = None, + max_seqlen: int | None = None, + **_: Any, +) -> torch.Tensor: + """Triton context attention without a causal mask.""" + cu_seqlens, seq_lens, max_seqlen = _varlen_metadata( + cu_seqlens, bsz, seq_len, device=q.device, max_seqlen=max_seqlen + ) + output = torch.empty_like(q) + context_attention_fwd( + q, + k, + v, + output, + cu_seqlens, + seq_lens, + max_seqlen, + is_causal=False, + sm_scale=softmax_scale, + ) + return output + + +def vision_attn_fa3( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + cu_seqlens: torch.Tensor | None, + bsz: int, + seq_len: int, + softmax_scale: float | None = None, + max_seqlen: int | None = None, + **_: Any, +) -> torch.Tensor: + cu_seqlens, _, max_seqlen = _varlen_metadata( + cu_seqlens, bsz, seq_len, device=q.device, max_seqlen=max_seqlen + ) + return 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, + softmax_scale=softmax_scale, + ) + + +def vision_attn_fa4( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + cu_seqlens: torch.Tensor | None, + bsz: int, + seq_len: int, + softmax_scale: float | None = None, + max_seqlen: int | None = None, + **_: Any, +) -> torch.Tensor: + cu_seqlens, _, max_seqlen = _varlen_metadata( + cu_seqlens, bsz, seq_len, device=q.device, max_seqlen=max_seqlen + ) + result = 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, + softmax_scale=softmax_scale, + ) + # FA4 CUTE returns (output, lse) in newer builds and bare output in older + # ones; downstream callers only consume the tensor. + return result[0] if isinstance(result, tuple) else result + + +def vision_attn_flashinfer_cudnn( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + cu_seqlens: torch.Tensor | None, + softmax_scale: float | None = None, + max_seqlen: Any = None, + sequence_lengths: torch.Tensor | None = None, + workspace_buffer: torch.Tensor | None = None, + **_: Any, +) -> torch.Tensor: + """cuDNN prefill backend. The caller (vision tower with cuDNN graph capture) + prepares ``cu_seqlens`` as three concatenated element-offset indptrs + ``[qk | v | o]`` of length ``batch+1`` each, plus ``sequence_lengths`` per + real (un-padded) sequence and ``max_seqlen`` as the bucketed budget. + """ + if ( + sequence_lengths is None + or max_seqlen is None + or not isinstance(cu_seqlens, torch.Tensor) + ): + raise ValueError( + "flashinfer_cudnn needs sequence_lengths, max_seqlen, and packed indptrs" + ) + + # cuDNN wants a python int for the seq budget. + max_seqlen = int( + max_seqlen.item() if isinstance(max_seqlen, torch.Tensor) else max_seqlen + ) + + # Flatten (b, s, h, d) -> (b*s, h, d) when the caller hands us 4-D. + in_4d = q.dim() == 4 + if in_4d: + b4 = q.shape[0] + q, k, v = (rearrange(t, "b s ... -> (b s) ...") for t in (q, k, v)) + + seq_lens = sequence_lengths.view(-1).to(device=q.device, dtype=torch.int32) + batch = seq_lens.numel() + packed = cu_seqlens.view(-1).to(device=q.device, dtype=torch.int32) + expected_packed = 3 * (batch + 1) + if packed.numel() != expected_packed: + raise ValueError( + f"Expected packed indptr length {expected_packed}, got {packed.numel()}." + ) + chunk = batch + 1 + qk_off = packed[:chunk].view(chunk, 1, 1, 1) + v_off = packed[chunk : 2 * chunk].view(chunk, 1, 1, 1) + o_off = packed[2 * chunk :].view(chunk, 1, 1, 1) + seq_lens_4d = seq_lens.view(batch, 1, 1, 1) + + head_size = q.shape[-1] + scale = softmax_scale if softmax_scale is not None else head_size**-0.5 + + output, _ = cudnn_batch_prefill_with_kv_cache( + q, + k, + v, + scale, + workspace_buffer, + max_token_per_sequence=max_seqlen, + max_sequence_kv=max_seqlen, + actual_seq_lens_q=seq_lens_4d, + actual_seq_lens_kv=seq_lens_4d, + causal=False, + return_lse=True, + batch_offsets_q=qk_off, + batch_offsets_k=qk_off, + batch_offsets_v=v_off, + batch_offsets_o=o_off, + is_cuda_graph_compatible=True, + ) + if in_4d: + output = rearrange(output, "(b s) h d -> b s h d", b=b4) + return output + + +_BACKENDS: dict[str, Callable[..., torch.Tensor]] = { + "triton_attn": vision_attn_triton, + "fa3": vision_attn_fa3, + "fa4": vision_attn_fa4, + "flashinfer_cudnn": vision_attn_flashinfer_cudnn, +} + + +def _default_multimodal_encoder_attn_backend() -> str: + """Platform default backend name.""" + if _is_nvidia: + if _platform.arch_version.major == 9: # Hopper SM90 + return "fa3" + if _platform.arch_version.major == 10: # Blackwell SM100 + return "fa4" + return "triton_attn" + if _is_amd: + return "triton_attn" + raise RuntimeError( + f"No default multimodal encoder attention backend for platform {_platform}; " + "set --mm-attention-backend explicitly." + ) + + +@functools.lru_cache(maxsize=None) +def _resolve_backend(name: str | None) -> Callable[..., torch.Tensor]: + """Resolve a backend name to its dispatch function. + + ``None`` falls back to the platform default; an unknown or platform- + incompatible name raises ValueError listing the registered backends. + Cached so a process logs the chosen backend exactly once per name. + """ + explicit = name is not None + if name is None: + name = _default_multimodal_encoder_attn_backend() + fn = _BACKENDS.get(name) + if fn is None: + raise ValueError( + f"Unknown multimodal encoder attention backend {name!r} " + f"(check --mm-attention-backend); available: {sorted(_BACKENDS)}" + ) + if name in ("fa3", "fa4", "flashinfer_cudnn") and not _is_nvidia: + raise ValueError( + f"multimodal encoder attention backend {name!r} is only available " + "on NVIDIA CUDA" + ) + if name == "fa3" and _platform.is_blackwell: + raise ValueError("The 'fa3' backend is not supported on Blackwell GPUs") + logger.info( + f"multimodal encoder attention backend: {name} " + f"({'override' if explicit else 'auto'})" + ) + return fn + + +class MultimodalEncoderAttention(nn.Module): + r"""Multi-headed attention without a KV cache for multimodal encoders.""" + + def __init__( + self, + embed_dim: int, + num_heads: int, + mapping: Mapping, + head_size: int | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + proj_bias: bool = True, + qkv_bias: bool = True, + customized_position_embedding_applier: Callable[ + [torch.Tensor, torch.Tensor, Any, Any], tuple[torch.Tensor, torch.Tensor] + ] = None, + position_embedding_mode: str | None = None, + workspace_buffer: torch.Tensor | None = None, + mm_attention_backend: str | None = None, + ): + super().__init__() + self.vision = mapping.vision + self.tp_size = self.vision.tp_size + self.tp_rank = self.vision.tp_rank + self.tp_group = self.vision.tp_group + self.head_size = head_size if head_size is not None else embed_dim // num_heads + self.num_attention_heads_per_partition = dist_utils.divide( + num_heads, self.tp_size + ) + self.num_attention_kv_heads_per_partition = dist_utils.divide( + 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.customized_position_embedding_applier = ( + customized_position_embedding_applier + ) + if position_embedding_mode not in (None, "complex_rope"): + raise ValueError( + f"Unknown vision position embedding mode: {position_embedding_mode}" + ) + self.position_embedding_mode = position_embedding_mode + self._backend_fn = _resolve_backend(mm_attention_backend) + self._use_packed_qkv_complex_rotary = ( + self._backend_fn is vision_attn_fa4 + and self.position_embedding_mode == "complex_rope" + ) + self._use_packed_qkv_rotary = ( + self._backend_fn is vision_attn_fa4 + and self.customized_position_embedding_applier is None + ) + self._copy_v_after_packed_qkv_rotary = False + self._workspace_buffer = workspace_buffer + + self.qkv_proj = QKVParallelLinear( + hidden_size=embed_dim, + head_size=self.head_size, + total_num_heads=num_heads, + total_num_kv_heads=num_heads, + bias=qkv_bias, + quant_config=quant_config, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + tp_group=self.tp_group, + prefix=add_prefix("qkv_proj", prefix), + ) + self.proj = RowParallelLinear( + input_size=num_heads * self.head_size, + output_size=embed_dim, + bias=proj_bias, + quant_config=quant_config, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + tp_group=self.tp_group, + prefix=add_prefix("proj", prefix), + reduce_results=True, + ) + + def forward( + self, + x: torch.Tensor, + cu_seqlens: torch.Tensor | None = None, + position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, + rotary_pos_emb_cos: torch.Tensor | None = None, + rotary_pos_emb_sin: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + r""" + Args: + x: [b, s, embed_dim] + cu_seqlens: [b] + Returns: + [b, s, head * head_size] + """ + if x.dim() == 2: + x = x.unsqueeze(0) + x_shape = x.shape + bsz, s, _ = x_shape + head = self.num_attention_heads_per_partition + kv_head = self.num_attention_kv_heads_per_partition + + max_seqlen = kwargs["max_seqlen"] if "max_seqlen" in kwargs else None + sequence_lengths = ( + kwargs["sequence_lengths"] if "sequence_lengths" in kwargs else None + ) + + qkv, _ = self.qkv_proj(x) + + use_packed_qkv_rotary = ( + self._use_packed_qkv_rotary + and position_embeddings is None + and rotary_pos_emb_cos is not None + and rotary_pos_emb_sin is not None + ) + use_packed_qkv_complex_rotary = ( + self._use_packed_qkv_complex_rotary and position_embeddings is not None + ) + cos = rotary_pos_emb_cos if use_packed_qkv_rotary else None + sin = rotary_pos_emb_sin if use_packed_qkv_rotary else None + + if use_packed_qkv_rotary: + if cos.size(-1) * 2 == self.head_size: + cos = torch.cat([cos, cos], dim=-1) + sin = torch.cat([sin, sin], dim=-1) + q, k, v = packed_qkv_neox_rotary( + qkv, + self.q_size, + self.kv_size, + head, + self.head_size, + cos, + sin, + copy_v=self._copy_v_after_packed_qkv_rotary, + ) + elif use_packed_qkv_complex_rotary: + q, k, v = packed_qkv_complex_rotary( + qkv, + self.q_size, + self.kv_size, + head, + self.head_size, + position_embeddings, + copy_v=self._copy_v_after_packed_qkv_rotary, + ) + else: + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + q = q.reshape(bsz * s, head, -1) + k = k.reshape(bsz * s, kv_head, -1) + v = v.reshape(bsz * s, kv_head, -1) + + cos = None + sin = None + + if position_embeddings is not None: + if self.customized_position_embedding_applier is not None: + q, k = self.customized_position_embedding_applier( + q, k, position_embeddings, x_shape + ) + else: + cos, sin = position_embeddings + elif rotary_pos_emb_cos is not None and rotary_pos_emb_sin is not None: + cos = rotary_pos_emb_cos + sin = rotary_pos_emb_sin + + if ( + not use_packed_qkv_rotary + and not use_packed_qkv_complex_rotary + and cos is not None + and sin is not None + ): + original_shape = q.shape + + # [total_tokens, head, head_size] + q = q.view(-1, head, self.head_size) + k = k.view(-1, head, self.head_size) + + if cos.size(-1) * 2 == self.head_size: + cos = torch.cat([cos, cos], dim=-1) + sin = torch.cat([sin, sin], dim=-1) + + q, k = apply_rotary_pos_emb_native(q, k, cos, sin) + q = q.view(original_shape) + k = k.view(original_shape) + + q, k, v = [ + rearrange(t, "b s ... -> (b s) ...") if t.dim() == 4 else t + for t in (q, k, v) + ] + + output = self._backend_fn( + q, + k, + v, + cu_seqlens=cu_seqlens, + bsz=bsz, + seq_len=s, + max_seqlen=max_seqlen, + sequence_lengths=sequence_lengths, + workspace_buffer=self._workspace_buffer, + ) + + output = rearrange(output, "(b s) ... h d -> b s ... (h d)", b=bsz) + output, _ = self.proj(output) + + return output + + +# Compatibility alias for existing vision tower implementations. +VisionAttention = MultimodalEncoderAttention diff --git a/python/tokenspeed/runtime/layers/attention/registry.py b/python/tokenspeed/runtime/layers/attention/registry.py new file mode 100644 index 0000000..cd2d06e --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/registry.py @@ -0,0 +1,813 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +from tokenspeed.runtime.configs.flat_memory_plan import ( + components_from_layers, + equalized_block_size, + plan_component_tensors, + state_const_bytes, +) +from tokenspeed.runtime.configs.model_config import AttentionArch, is_deepseek_v4 +from tokenspeed.runtime.configs.paged_cache_spec import ( + hybrid_slab_group_size, + scheduler_ext_flat_kvcache, +) +from tokenspeed.runtime.layers.attention.configs.base import BaseAttnConfig +from tokenspeed.runtime.layers.attention.configs.dsa import DSAConfig +from tokenspeed.runtime.layers.attention.configs.mha import MHAConfig +from tokenspeed.runtime.layers.attention.configs.mla import MLAConfig +from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool +from tokenspeed.runtime.layers.attention.utils import ( + profile_available_cache_memory_bytes, + profile_cache_budget, + profile_max_num_pages, +) +from tokenspeed.runtime.utils.env import envs + +logger = logging.getLogger(__name__) + +_CI_SMALL_KV_SIZE = envs.TOKENSPEED_CI_SMALL_KV_SIZE.get_set_value_or(None) + +if TYPE_CHECKING: + from tokenspeed.runtime.configs.model_config import ModelConfig + from tokenspeed.runtime.layers.attention.backends.base import AttentionBackend + from tokenspeed.runtime.utils.server_args import ServerArgs + + +def _kv_profile_layer_divisor(num_layers, layer_types, *, sliding_window_tokens=None): + """Attention layers to charge per token in the KV memory profile: + layers-per-group under the slab layout, else all layers (single + source: hybrid_slab_group_size).""" + gs = hybrid_slab_group_size( + layer_types, + sliding_window_tokens=sliding_window_tokens, + ) + return gs if gs is not None else num_layers + + +def _resolve_max_num_tokens( + profiled_num_pages: int, + page_size: int, + max_total_tokens: int | None, +) -> int: + profiled_tokens = profiled_num_pages * page_size + if max_total_tokens is None: + return profiled_tokens + requested_pages = max_total_tokens // page_size + if requested_pages < 1: + raise ValueError( + f"max_total_tokens={max_total_tokens} must contain at least one full page " + f"(page_size={page_size})" + ) + return min(profiled_tokens, requested_pages * page_size) + + +def _resolve_draft_cache_cell_size_for_profile( + draft_attn_config: BaseAttnConfig | None, + draft_model_config: ModelConfig | None, + draft_profile_cache_cell_size: int | None, +) -> int: + if draft_profile_cache_cell_size is not None: + return draft_profile_cache_cell_size + if draft_attn_config is None or draft_model_config is None: + return 0 + return draft_attn_config.cache_cell_size() * draft_model_config.num_attention_layers + + +# ---------- backend registry ---------- + +# Maps backend_name -> (supported archs, backend class) +_BACKEND_REGISTRY: dict[str, tuple[set[AttentionArch], type[AttentionBackend]]] = {} + + +def register_backend( + name: str, + archs: set[AttentionArch], + cls: type[AttentionBackend], +) -> None: + _BACKEND_REGISTRY[name] = (archs, cls) + + +_HYBRID_GDN_ARCHITECTURES = { + "Qwen3_5MoeForConditionalGeneration", + "Qwen3_5MoeForConditionalGenerationNextN", + "Qwen3_5ForConditionalGeneration", + "Qwen3_5ForConditionalGenerationNextN", +} + + +# Aliases for backward compatibility with server_args choices +_BACKEND_ALIASES = { + "trtllm_mha": "trtllm", +} + + +def _get_default_backend_name(arch: AttentionArch) -> str: + if arch == AttentionArch.MLA: + return "mla" + if arch == AttentionArch.DSA: + return "dsa" + else: + return "mha" + + +def _get_backend_cls(name: str, arch: AttentionArch) -> type[AttentionBackend]: + if name is None: + candidates = [_get_default_backend_name(arch)] + for candidate in candidates: + entry = _BACKEND_REGISTRY.get(candidate) + if entry is not None and arch in entry[0]: + return entry[1] + raise ValueError( + f"No backend supports arch {arch}. Available: {list(_BACKEND_REGISTRY)}" + ) + name = _BACKEND_ALIASES.get(name, name) + entry = _BACKEND_REGISTRY.get(name) + if entry is None: + raise ValueError( + f"Unknown attention backend: {name!r}. Available: {list(_BACKEND_REGISTRY)}" + ) + supported_archs, cls = entry + if arch not in supported_archs: + raise ValueError( + f"Backend {name!r} does not support arch {arch}. " + f"Supported archs: {supported_archs}" + ) + return cls + + +# ---------- arch -> config class ---------- + +_CONFIG_CLS: dict[AttentionArch, type[BaseAttnConfig]] = { + AttentionArch.MHA: MHAConfig, + AttentionArch.MLA: MLAConfig, + AttentionArch.DSA: DSAConfig, +} + + +def _create_attn_config( + server_args: ServerArgs, model_config: ModelConfig, is_draft: bool = False +) -> BaseAttnConfig: + arch = model_config.attention_arch + if arch not in _CONFIG_CLS: + raise NotImplementedError(f"Not supported Attention Arch: {arch!r}") + return _CONFIG_CLS[arch].generate(server_args, model_config, is_draft) + + +def _create_attn_backend( + arch: AttentionArch, + config: BaseAttnConfig, +) -> AttentionBackend: + return _get_backend_cls(config.backend_name, arch)(config) + + +def _create_attn_backend_with_name( + name: str | None, + arch: AttentionArch, + config: BaseAttnConfig, +) -> AttentionBackend: + original_name = config.backend_name + config.backend_name = name + try: + return _get_backend_cls(name, arch)(config) + finally: + config.backend_name = original_name + + +def _create_attn_pool( + config: BaseAttnConfig, + num_layers: int, + max_total_num_tokens: int, + rank: int, + enable_memory_saver: bool = False, +) -> BaseTokenToKVPool: + return config.create_pool( + num_layers, max_total_num_tokens, rank, enable_memory_saver + ) + + +def _attention_use_fp4_indexer_cache(server_args: "ServerArgs", hf_config) -> bool: + if getattr(server_args, "attention_use_fp4_indexer_cache", None) is not None: + return bool(server_args.attention_use_fp4_indexer_cache) + attention_config = getattr(hf_config, "attention_config", None) + if isinstance(attention_config, dict): + return bool(attention_config.get("use_fp4_indexer_cache", False)) + return bool(getattr(attention_config, "use_fp4_indexer_cache", False)) + + +def _create_hybrid_linear_attn( + server_args: ServerArgs, + model_config: ModelConfig, + config: BaseAttnConfig, + arch: AttentionArch, + max_num_tokens: int, + rank: int, + enable_memory_saver: bool = False, + full_attn_backend_name: str = None, + mamba_pool_total_chunks: int = 0, +) -> tuple[AttentionBackend, BaseTokenToKVPool, object]: + """Create a hybrid backend + pool for GDN hybrid models (Qwen3.5, Qwen3Next).""" + from tokenspeed.runtime.layers.attention.backends.hybrid_linear_attn import ( + HybridLinearAttnBackend, + LayerMappedKVPool, + MambaAttnBackend, + SimpleMambaPool, + ) + + hf_config = model_config.hf_config + text_config = getattr(hf_config, "text_config", hf_config) + full_attn_layers = text_config.full_attention_layer_ids + + # Create the full attention backend for standard MHA layers. + # Use user's original choice if provided, otherwise auto-select. + full_attn_backend = _create_attn_backend_with_name( + full_attn_backend_name, + arch, + config, + ) + + # Create mamba/linear attention backend. Only propagate the configured + # verify width when spec-dec is actually enabled — matches MLAConfig / + # MHAConfig.generate. Otherwise the BaseAttnConfig sentinel (1) wins so + # non-spec hybrid decode doesn't get misclassified as target verify / + # draft extend by `self.spec_num_tokens > 1`. + if server_args.speculative_algorithm is not None: + config.speculative_num_draft_tokens = server_args.speculative_num_draft_tokens + + flat_kvcache = scheduler_ext_flat_kvcache() + if flat_kvcache: + # Flat path: the pool covers ALL layers, so pool indices == global + # layer ids, its layer_types line up with the state slabs, and the + # group specs publish both the "full_attention" and + # "linear_attention" groups. State layers carry NO k/v tensors + # (None slots, M18a T4) -- matching the plan sizing, which charges + # only full-layer KV + state rows. The identity mapping keeps + # the wrapper type identical to the radix path. + num_total_layers = len(text_config.layers_block_type) + inner_pool = config.create_pool( + num_total_layers, max_num_tokens, rank, enable_memory_saver + ) + pool = LayerMappedKVPool(inner_pool, list(range(num_total_layers))) + else: + # Create KV cache pool (only for full attention layers) + num_full_attn_layers = len(full_attn_layers) + inner_pool = config.create_pool( + num_full_attn_layers, max_num_tokens, rank, enable_memory_saver + ) + # Wrap with layer ID mapping (global layer IDs -> pool indices) + pool = LayerMappedKVPool(inner_pool, full_attn_layers) + + # Read mamba2_cache_params to decide whether this model actually has + # any linear / mamba layers. A draft model on a hybrid-GDN target + # (e.g. MTP on Qwen3.5) shares the same architecture class as the + # target but commonly ships with *zero* mamba layers — in that case + # we skip the mamba backend / pool entirely so that its + # ``init_forward_metadata_*`` hooks do not run (they would otherwise + # touch a zero-sized pool on the same persistent state_indices_list + # as the target, which breaks the captured CUDA graph). + ( + conv_state_shape, + temporal_state_shape, + conv_dtype, + ssm_dtype, + mamba_layer_ids, + ) = text_config.mamba2_cache_params + + if len(mamba_layer_ids) == 0: + logger.info( + "Created hybrid_linear_attn backend: %d full attn layers, 0 linear " + "attn layers (skipping mamba backend / pool)", + len(full_attn_layers), + ) + return full_attn_backend, pool, None + + linear_attn_backend = MambaAttnBackend(config) + + if flat_kvcache: + # Flat mode never touches a SimpleMambaPool: the recurrent state + # lives in the KV pool's state slabs, addressed by the flat block + # tables (set_kv_pool below activates the dual-index state paging), + # so skip the pool and its set_pool binding entirely. + mamba_pool = None + else: + # Mamba radix cache uses C++ chunk indices. Without radix cache, the + # backend uses 1-based req_pool_indices directly, so keep slot 0 as + # padding. + per_rank_max_batch = server_args.max_num_seqs // max( + server_args.data_parallel_size or server_args.mapping.attn.dp_size, 1 + ) + req_pool_padding_index = per_rank_max_batch + 1 + mamba_pool_size = ( + mamba_pool_total_chunks + 1 + if mamba_pool_total_chunks > 0 + else per_rank_max_batch + 1 + ) + mamba_pool = SimpleMambaPool( + size=mamba_pool_size, + num_mamba_layers=len(mamba_layer_ids), + conv_state_shape=conv_state_shape, + temporal_state_shape=temporal_state_shape, + conv_dtype=conv_dtype, + ssm_dtype=ssm_dtype, + mamba_layer_ids=mamba_layer_ids, + device=config.device, + page_size=server_args.block_size, + speculative_num_draft_tokens=( + server_args.speculative_num_draft_tokens + if server_args.speculative_algorithm is not None + else 0 + ), + # ``current_input_indices`` is keyed by the scheduler's rank-local, + # 1-based req_pool_idx; the row after that range is the CUDA graph + # padding sentinel. + max_req_pool_size=req_pool_padding_index, + ) + linear_attn_backend.set_pool(mamba_pool) + # Flat state paging (dual-index) keys off the KV pool's state slabs + + # published "linear_attention" group; no-op on the radix path. + linear_attn_backend.set_kv_pool(pool) + + backend = HybridLinearAttnBackend( + full_attn_backend, linear_attn_backend, full_attn_layers + ) + logger.info( + "Created hybrid_linear_attn backend: %d full attn layers, %d linear attn layers, %s", + len(full_attn_layers), + len(mamba_layer_ids), + ( + "flat state slabs (no mamba slot pool)" + if mamba_pool is None + else f"mamba pool size {mamba_pool.size}" + ), + ) + return backend, pool, mamba_pool + + +# ---------- public API ---------- +def create_attn_components( + server_args: ServerArgs, + model_config: ModelConfig, + gpu_id: int, + rank: int, + gpu_memory: int, + enable_memory_saver: bool = False, + draft_model_config: ModelConfig | None = None, + decode_input_tokens: int = 1, + overlap_schedule_depth: int = 0, +) -> tuple[ + AttentionBackend, + BaseTokenToKVPool, + AttentionBackend | None, + BaseTokenToKVPool | None, + int, + int, + object | None, +]: + arch = model_config.attention_arch + + architectures = getattr(model_config.hf_config, "architectures", None) or [] + is_hybrid_gdn = any(a in _HYBRID_GDN_ARCHITECTURES for a in architectures) + is_deepseek_v4_model = is_deepseek_v4(model_config.hf_config) + is_deepseek_v4_draft_model = draft_model_config is not None and is_deepseek_v4( + draft_model_config.hf_config + ) + original_attn_backend = server_args.attention_backend + if is_deepseek_v4_model: + server_args.attention_backend = "deepseek_v4" + if is_deepseek_v4_draft_model: + server_args.drafter_attention_backend = "deepseek_v4" + if is_hybrid_gdn: + # Qwen3.5 GDN hybrid models always need hybrid_linear_attn. + # Save the user's original choice for the full-attention sub-backend. + server_args.attention_backend = "hybrid_linear_attn" + elif server_args.attention_backend == "hybrid_linear_attn": + logger.warning( + "Ignoring hybrid_linear_attn backend for non-hybrid model architectures=%s", + architectures, + ) + server_args.attention_backend = None + if server_args.drafter_attention_backend == "hybrid_linear_attn": + server_args.drafter_attention_backend = None + + config = _create_attn_config(server_args, model_config) + is_flat_gdn = getattr(config, "conv_state_shape", None) is not None + gdn_state_bytes = ( + state_const_bytes( + config.conv_state_shape, + config.conv_dtype, + config.temporal_state_shape, + config.ssm_dtype, + ) + if is_flat_gdn + else None + ) + if is_flat_gdn: + equalized_block_size_value = equalized_block_size( + layer_types=list(config.layer_types), + kv_bytes_per_slot=config.cache_cell_size(), + state_const_bytes=gdn_state_bytes, + block_size=server_args.block_size, + ) + if equalized_block_size_value != server_args.block_size: + logger.info( + "Setting attention block size to %d tokens to cover the GDN " + "state row (configured block size %d)", + equalized_block_size_value, + server_args.block_size, + ) + server_args.block_size = equalized_block_size_value + config.page_size = equalized_block_size_value + draft_attn_config = None + if draft_model_config: + draft_attn_config = _create_attn_config( + server_args, draft_model_config, is_draft=True + ) + num_layers = model_config.num_attention_layers + deepseek_v4_layout = None + draft_deepseek_v4_layout = None + profile_cache_cell_size = None + draft_profile_cache_cell_size = None + if is_deepseek_v4_model: + from tokenspeed.runtime.layers.attention.kv_cache.deepseek_v4 import ( + deepseek_v4_cache_layout_from_config, + ) + + deepseek_v4_layout = deepseek_v4_cache_layout_from_config( + model_config.hf_config, + page_size=server_args.block_size, + use_fp4_indexer_cache=_attention_use_fp4_indexer_cache( + server_args, model_config.hf_config + ), + layer_indices=range(num_layers), + ) + profile_cache_cell_size = deepseek_v4_layout.cache_cell_size(num_layers) + if is_deepseek_v4_draft_model: + from tokenspeed.runtime.layers.attention.kv_cache.deepseek_v4 import ( + deepseek_v4_cache_layout_from_config, + ) + + draft_layer_start = draft_model_config.num_hidden_layers + draft_num_layers = draft_model_config.num_attention_layers + draft_deepseek_v4_layout = deepseek_v4_cache_layout_from_config( + draft_model_config.hf_config, + page_size=server_args.block_size, + use_fp4_indexer_cache=_attention_use_fp4_indexer_cache( + server_args, draft_model_config.hf_config + ), + layer_indices=range( + draft_layer_start, + draft_layer_start + draft_num_layers, + ), + ) + draft_profile_cache_cell_size = draft_deepseek_v4_layout.cache_cell_size( + draft_model_config.num_attention_layers + ) + + hf_config = getattr(model_config, "hf_config", None) + text_config = getattr(hf_config, "text_config", hf_config) if hf_config else None + mamba_cache_params = ( + getattr(text_config, "mamba2_cache_params", None) if text_config else None + ) + # Unpack once with names; every consumer below reads these instead of + # indexing into the raw tuple. + if mamba_cache_params: + ( + mamba_conv_state_shape, + mamba_temporal_state_shape, + mamba_conv_dtype, + mamba_ssm_dtype, + mamba_layer_ids, + ) = mamba_cache_params + else: + mamba_conv_state_shape = mamba_temporal_state_shape = None + mamba_conv_dtype = mamba_ssm_dtype = None + mamba_layer_ids = () + has_mamba_layers = len(mamba_layer_ids) > 0 + has_mamba = getattr(model_config, "mambaish_config", None) is not None or ( + has_mamba_layers + ) + mamba_pool_total_chunks = 0 + mamba_pool = None + + _profile_kwargs = dict( + attn_config=config, + gpu_id=gpu_id, + tp_size=server_args.mapping.world_size, + page_size=server_args.block_size, + num_attention_layers=num_layers, + total_gpu_memory=gpu_memory, + world_group=server_args.mapping.world_group, + draft_attn_config=draft_attn_config if draft_attn_config else None, + draft_num_attention_layers=( + draft_model_config.num_attention_layers if draft_attn_config else None + ), + ) + + if is_deepseek_v4_model: + from tokenspeed.runtime.layers.attention.kv_cache.deepseek_v4 import ( + profile_deepseek_v4_max_num_pages, + ) + + draft_cache_cell_size = _resolve_draft_cache_cell_size_for_profile( + draft_attn_config, + draft_model_config, + draft_profile_cache_cell_size, + ) + max_total_num_pages = profile_deepseek_v4_max_num_pages( + layout=deepseek_v4_layout, + hf_config=model_config.hf_config, + layer_num=num_layers, + max_live_requests=config.max_bs, + max_scheduled_tokens=server_args.chunked_prefill_size, + max_context_len=config.context_len, + available_cache_memory_bytes=profile_available_cache_memory_bytes( + attn_config=config, + gpu_id=gpu_id, + tp_size=server_args.mapping.world_size, + gpu_memory_utilization=server_args.gpu_memory_utilization, + total_gpu_memory=gpu_memory, + world_group=server_args.mapping.world_group, + ), + draft_cache_cell_size=draft_cache_cell_size, + decode_input_tokens=decode_input_tokens, + overlap_schedule_depth=overlap_schedule_depth, + ) + logger.info( + "DeepSeek V4 grouped KV profile: max_live_requests=%s " + "(attn config max_bs=%s, attn_dp_size=%s), max_total_num_pages=%s", + config.max_bs, + config.max_bs, + server_args.mapping.attn.dp_size, + max_total_num_pages, + ) + max_num_tokens = _resolve_max_num_tokens( + max_total_num_pages, + server_args.block_size, + server_args.max_total_tokens, + ) + elif has_mamba and is_flat_gdn: + draft_row_bytes = 0 + if draft_attn_config is not None: + draft_row_bytes = ( + _resolve_draft_cache_cell_size_for_profile( + draft_attn_config, draft_model_config, draft_profile_cache_cell_size + ) + * server_args.block_size + ) + cache_memory = profile_available_cache_memory_bytes( + attn_config=config, + gpu_id=gpu_id, + tp_size=server_args.mapping.world_size, + gpu_memory_utilization=server_args.gpu_memory_utilization, + total_gpu_memory=gpu_memory, + world_group=server_args.mapping.world_group, + ) + flat_plan = plan_component_tensors( + components_from_layers( + layer_types=list(config.layer_types), + kv_bytes_per_slot=config.cache_cell_size(), + state_const_bytes=gdn_state_bytes, + ), + block_size=server_args.block_size, + budget_bytes=cache_memory, + reserved_bytes_per_block=draft_row_bytes, + ) + max_total_num_pages = flat_plan.geometry.num_blocks + logger.info( + "Flat GDN KV profile: block_bytes=%d (%d component tensors, " + "block_size=%d), max_total_num_pages=%d", + flat_plan.geometry.block_bytes, + len(flat_plan.tensors), + server_args.block_size, + max_total_num_pages, + ) + max_num_tokens = _resolve_max_num_tokens( + max_total_num_pages, server_args.block_size, server_args.max_total_tokens + ) + elif has_mamba and server_args.max_mamba_cache_size is not None: + mamba_pool_total_chunks = server_args.max_mamba_cache_size + full_attn_layer_ids = getattr(text_config, "full_attention_layer_ids", None) + num_kv_layers = ( + len(full_attn_layer_ids) + if full_attn_layer_ids is not None + else num_layers - len(mamba_layer_ids) + ) + max_total_num_pages = profile_max_num_pages( + **{**_profile_kwargs, "num_attention_layers": num_kv_layers}, + gpu_memory_utilization=server_args.gpu_memory_utilization, + cache_cell_size=profile_cache_cell_size, + draft_cache_cell_size=draft_profile_cache_cell_size, + ) + max_num_tokens = _resolve_max_num_tokens( + max_total_num_pages, + server_args.block_size, + server_args.max_total_tokens, + ) + elif has_mamba and server_args.max_mamba_cache_size is None: + num_mamba_layers = len(mamba_layer_ids) + speculative_num_draft_tokens = ( + server_args.speculative_num_draft_tokens + if server_args.speculative_algorithm is not None + else 0 + ) + per_layer_mamba_chunk_memory = sum( + state_const_bytes( + mamba_conv_state_shape, + mamba_conv_dtype, + mamba_temporal_state_shape, + mamba_ssm_dtype, + ).values() + ) * (1 + speculative_num_draft_tokens) + memory_per_mamba_chunk = num_mamba_layers * per_layer_mamba_chunk_memory + full_attn_layer_ids = getattr(text_config, "full_attention_layer_ids", None) + num_kv_layers = ( + len(full_attn_layer_ids) + if full_attn_layer_ids is not None + else num_layers - num_mamba_layers + ) + kv_max_num_pages, mamba_pool_total_chunks = profile_cache_budget( + **{**_profile_kwargs, "num_attention_layers": num_kv_layers}, + mem_fraction_static=server_args.gpu_memory_utilization, + mamba_memory_per_chunk=memory_per_mamba_chunk, + mamba_ratio=server_args.mamba_full_memory_ratio, + ) + max_num_tokens = _resolve_max_num_tokens( + kv_max_num_pages, + server_args.block_size, + server_args.max_total_tokens, + ) + else: + # config.layer_types / config.sliding_window_tokens are the exact + # values forwarded to the KV pool, so sizing and layout consume + # identical inputs (MLA configs carry neither -> legacy divisor). + slab_divisor = _kv_profile_layer_divisor( + num_layers, + getattr(config, "layer_types", None), + sliding_window_tokens=getattr(config, "sliding_window_tokens", None), + ) + if profile_cache_cell_size is not None and slab_divisor != num_layers: + # A cell-size override can't compose with the slab divisor. + logger.warning( + "hybrid slab sizing disabled: profile cache_cell_size " + "override is set; charging all %d layers instead of %d", + num_layers, + slab_divisor, + ) + slab_divisor = num_layers + max_total_num_pages = profile_max_num_pages( + **{**_profile_kwargs, "num_attention_layers": slab_divisor}, + gpu_memory_utilization=server_args.gpu_memory_utilization, + cache_cell_size=profile_cache_cell_size, + draft_cache_cell_size=draft_profile_cache_cell_size, + ) + max_num_tokens = _resolve_max_num_tokens( + max_total_num_pages, + server_args.block_size, + server_args.max_total_tokens, + ) + + if _CI_SMALL_KV_SIZE is not None and int(_CI_SMALL_KV_SIZE) > 0: + max_num_tokens = int(_CI_SMALL_KV_SIZE) + if max_num_tokens <= 0: + raise ValueError( + f"KV cache token pool size must be positive, got {max_num_tokens}" + ) + + if is_deepseek_v4_model: + from tokenspeed.runtime.layers.attention.kv_cache.deepseek_v4 import ( + DeepseekV4TokenToKVPool, + ) + + backend = _create_attn_backend(arch, config) + pool = DeepseekV4TokenToKVPool( + size=max_num_tokens, + model_dtype=model_config.dtype, + layout=deepseek_v4_layout, + layer_num=num_layers, + device=config.device, + enable_memory_saver=enable_memory_saver, + max_batch_size=config.max_bs, + max_context_len=config.context_len, + page_size=server_args.block_size, + rank=rank, + hf_config=model_config.hf_config, + max_scheduled_tokens=server_args.chunked_prefill_size, + decode_input_tokens=decode_input_tokens, + overlap_schedule_depth=overlap_schedule_depth, + ) + elif is_hybrid_gdn: + resolved_original_backend = _BACKEND_ALIASES.get( + original_attn_backend, original_attn_backend + ) + backend, pool, mamba_pool = _create_hybrid_linear_attn( + server_args, + model_config, + config, + arch, + max_num_tokens, + rank, + enable_memory_saver, + full_attn_backend_name=( + resolved_original_backend + if resolved_original_backend != "hybrid_linear_attn" + else None + ), + mamba_pool_total_chunks=mamba_pool_total_chunks, + ) + else: + backend = _create_attn_backend(arch, config) + pool = _create_attn_pool( + config, num_layers, max_num_tokens, rank, enable_memory_saver + ) + draft_attn_backend = None + draft_pool = None + if draft_attn_config: + # Check if draft model is also a hybrid GDN model. + draft_archs = getattr(draft_model_config.hf_config, "architectures", None) or [] + if is_deepseek_v4_draft_model: + from tokenspeed.runtime.layers.attention.kv_cache.deepseek_v4 import ( + DeepseekV4TokenToKVPool, + ) + + draft_attn_backend = _create_attn_backend( + draft_model_config.attention_arch, draft_attn_config + ) + draft_pool = DeepseekV4TokenToKVPool( + size=max_num_tokens, + model_dtype=draft_model_config.dtype, + layout=draft_deepseek_v4_layout, + layer_num=draft_model_config.num_attention_layers, + device=draft_attn_config.device, + enable_memory_saver=enable_memory_saver, + max_batch_size=draft_attn_config.max_bs, + max_context_len=draft_attn_config.context_len, + page_size=server_args.block_size, + rank=rank, + hf_config=draft_model_config.hf_config, + max_scheduled_tokens=server_args.chunked_prefill_size, + decode_input_tokens=decode_input_tokens, + overlap_schedule_depth=overlap_schedule_depth, + ) + elif any(a in _HYBRID_GDN_ARCHITECTURES for a in draft_archs): + resolved_draft_backend = _BACKEND_ALIASES.get( + original_attn_backend, original_attn_backend + ) + draft_attn_backend, draft_pool, _ = _create_hybrid_linear_attn( + server_args, + draft_model_config, + draft_attn_config, + draft_model_config.attention_arch, + max_num_tokens, + rank, + enable_memory_saver, + full_attn_backend_name=( + resolved_draft_backend + if resolved_draft_backend != "hybrid_linear_attn" + else None + ), + mamba_pool_total_chunks=mamba_pool_total_chunks, + ) + else: + draft_attn_backend = _create_attn_backend( + draft_model_config.attention_arch, draft_attn_config + ) + draft_layers = draft_model_config.num_attention_layers + draft_pool = _create_attn_pool( + draft_attn_config, + draft_layers, + max_num_tokens, + rank, + enable_memory_saver, + ) + + return ( + backend, + pool, + draft_attn_backend, + draft_pool, + max_num_tokens, + mamba_pool_total_chunks, + mamba_pool, + ) diff --git a/python/tokenspeed/runtime/layers/attention/utils.py b/python/tokenspeed/runtime/layers/attention/utils.py new file mode 100755 index 0000000..eab2cdd --- /dev/null +++ b/python/tokenspeed/runtime/layers/attention/utils.py @@ -0,0 +1,223 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import torch +import triton +import triton.language as tl + +from tokenspeed.runtime.distributed.process_group_manager import ( + process_group_manager as pg_manager, +) +from tokenspeed.runtime.layers.attention.configs.base import BaseAttnConfig +from tokenspeed.runtime.utils import get_available_gpu_memory + + +@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): + offset = tl.arange(0, BLOCK_SIZE) + 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) + + +# --- Page table helpers (shared across attention backends) --- + + +def build_page_table( + req_pool_indices: torch.Tensor, + req_to_page: torch.Tensor, + page_size: int, + max_seq_len_k: int, +) -> torch.Tensor: + """Build page table from req_to_page. + + req_to_page: [req_pool_size+1, max_pages] containing page IDs. + Returns: [bs, max_pages_needed] page table slice. + """ + max_pages = (max_seq_len_k + page_size - 1) // page_size + return req_to_page[req_pool_indices, :max_pages] + + +def update_page_table_inplace( + page_table_buf: torch.Tensor, + req_pool_indices: torch.Tensor, + req_to_page: torch.Tensor, + page_size: int, + max_seq_len_k: int, +): + """Copy page table from req_to_page into pre-allocated CUDA graph buffer.""" + max_pages = (max_seq_len_k + page_size - 1) // page_size + page_table_buf[:, :max_pages].copy_(req_to_page[req_pool_indices, :max_pages]) + + +def token_indices_from_pages( + req_pool_indices: torch.Tensor, + token_positions: torch.Tensor, + req_to_page: torch.Tensor, + page_size: int, +) -> torch.Tensor: + """Convert token positions to KV slot indices using req_to_page. + + token_positions: [bs, num_tokens] — token offsets within each request. + Returns: [bs, num_tokens] — KV cache slot IDs (page_id * page_size + offset). + """ + page_indices = token_positions // page_size + offsets = token_positions % page_size + page_ids = req_to_page[req_pool_indices].gather(1, page_indices) + return page_ids * page_size + offsets + + +# --- Page-based memory profiling --- + + +def profile_available_cache_memory_bytes( + attn_config: BaseAttnConfig, + gpu_id: int, + tp_size: int, + gpu_memory_utilization: float, + total_gpu_memory: int, + world_group=None, +) -> int: + cpu_group = ( + pg_manager.get_process_group("gloo", world_group) + if world_group is not None + else None + ) + available_gpu_memory = get_available_gpu_memory( + attn_config.device, + gpu_id, + distributed=tp_size > 1, + cpu_group=cpu_group, + ) + cache_memory = available_gpu_memory - total_gpu_memory * ( + 1 - gpu_memory_utilization + ) + return int(cache_memory * (1 << 30)) + + +def profile_max_num_pages( + attn_config: BaseAttnConfig, + gpu_id: int, + tp_size: int, + gpu_memory_utilization: float, + page_size: int, + num_attention_layers: int, + total_gpu_memory: int, + world_group=None, + draft_attn_config: BaseAttnConfig | None = None, + draft_num_attention_layers: int | None = None, + cache_cell_size: int | None = None, + draft_cache_cell_size: int | None = None, +): + cache_memory = profile_available_cache_memory_bytes( + attn_config, + gpu_id, + tp_size, + gpu_memory_utilization, + total_gpu_memory, + world_group, + ) + if cache_cell_size is None: + cell_size = attn_config.cache_cell_size() * num_attention_layers + else: + cell_size = cache_cell_size + if draft_attn_config is not None: + if draft_cache_cell_size is None: + cell_size += ( + draft_attn_config.cache_cell_size() * draft_num_attention_layers + ) + else: + cell_size += draft_cache_cell_size + if cell_size <= 0: + raise ValueError(f"KV cache cell size must be positive, got {cell_size}") + max_num_token = cache_memory // cell_size + max_num_pages = (max_num_token + page_size - 1) // page_size + return max_num_pages + + +def profile_cache_budget( + attn_config: BaseAttnConfig, + gpu_id: int, + tp_size: int, + mem_fraction_static: float, + page_size: int, + num_attention_layers: int, + total_gpu_memory: int, + mamba_memory_per_chunk: int, + mamba_ratio: float, + world_group=None, + draft_attn_config: BaseAttnConfig | None = None, + draft_num_attention_layers: int | None = None, +) -> tuple[int, int]: + """Profile GPU memory and split between KV pages and mamba chunks. + + Returns: + (kv_max_num_pages, mamba_pool_total_chunks) + """ + total_cache_memory = profile_available_cache_memory_bytes( + attn_config, + gpu_id, + tp_size, + mem_fraction_static, + total_gpu_memory, + world_group, + ) + cell_size = attn_config.cache_cell_size() * num_attention_layers + if draft_attn_config is not None: + cell_size += draft_attn_config.cache_cell_size() * draft_num_attention_layers + + kv_memory = int(total_cache_memory / (1 + mamba_ratio)) + mamba_memory = total_cache_memory - kv_memory + + kv_cell_size = cell_size * page_size + kv_max_num_pages = int(kv_memory // kv_cell_size) + mamba_pool_total_chunks = max(int(mamba_memory // mamba_memory_per_chunk), 2) + + return kv_max_num_pages, mamba_pool_total_chunks diff --git a/python/tokenspeed/runtime/layers/common.py b/python/tokenspeed/runtime/layers/common.py new file mode 100644 index 0000000..1c6bcb9 --- /dev/null +++ b/python/tokenspeed/runtime/layers/common.py @@ -0,0 +1,38 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Small compiled tensor helpers shared across runtime layers.""" + +from __future__ import annotations + +import torch +from tokenspeed_kernel.torch_compile import get_compiler_backend + + +@torch.compile(dynamic=True, backend=get_compiler_backend()) +def concat(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: + """Concatenate two tensors along the last dimension.""" + return torch.cat([a, b], dim=-1) + + +@torch.compile(dynamic=True, backend=get_compiler_backend()) +def fp8_cast_contiguous(x: torch.Tensor) -> torch.Tensor: + """Cast to FP8 E4M3 and return a contiguous tensor.""" + return x.to(torch.float8_e4m3fn).contiguous() diff --git a/python/tokenspeed/runtime/layers/conv.py b/python/tokenspeed/runtime/layers/conv.py new file mode 100644 index 0000000..2e0ba2d --- /dev/null +++ b/python/tokenspeed/runtime/layers/conv.py @@ -0,0 +1,312 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright 2023-2024 SGLang Team +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +""" +Conv2d/Conv3d layers with unfold+linear optimization for patch embeddings. + +When kernel_size == stride, padding == 0, dilation == 1, groups == 1, the conv +is equivalent to unfold + F.linear, which is significantly faster on CUDA and +also avoids the PyTorch 2.9.1 + CuDNN < 9.15 Conv3d bug +(https://github.com/pytorch/pytorch/issues/168167). +""" + +import math + +import torch +import torch.nn as nn +import torch.nn.functional as F + +_VALID_PADDING_STRINGS = {"same", "valid"} +_VALID_PADDING_MODES = {"zeros", "reflect", "replicate", "circular"} + + +def _tuplify(val, n: int) -> tuple: + if isinstance(val, (list, tuple)): + if len(val) != n: + raise ValueError(f"Expected {n} values, got {len(val)}.") + return tuple(val) + return (val,) * n + + +def _check_enable_linear( + kernel_size: tuple, + stride: tuple, + padding: tuple, + dilation: tuple, + groups: int, +) -> bool: + """Check if conv can be replaced with unfold + F.linear.""" + return ( + kernel_size == stride + and all(p == 0 for p in padding) + and all(d == 1 for d in dilation) + and groups == 1 + ) + + +def _reverse_repeat_tuple(t: tuple) -> tuple: + """(1, 2, 3) -> (3, 3, 2, 2, 1, 1). Used for F.pad with non-zeros padding_mode.""" + return tuple(x for x in reversed(t) for _ in range(2)) + + +def _compute_same_padding_for_pad(kernel_size: tuple, dilation: tuple) -> tuple: + """Compute _reversed_padding_repeated_twice for padding='same'. + + This mirrors PyTorch's nn.Conv*d behavior: pre-compute the exact pad + amounts so that F.pad can be called before F.conv*d(padding=0). + """ + pad = [] + for k, d in zip(reversed(kernel_size), reversed(dilation)): + total = d * (k - 1) + pad.append(total // 2) + pad.append(total - total // 2) + return tuple(pad) + + +def _validate_conv_args( + in_channels: int, + out_channels: int, + groups: int, + padding, + padding_mode: str, + stride: tuple, +) -> None: + if in_channels % groups != 0: + raise ValueError( + f"in_channels ({in_channels}) must be divisible by groups ({groups})" + ) + if out_channels % groups != 0: + raise ValueError( + f"out_channels ({out_channels}) must be divisible by groups ({groups})" + ) + if padding_mode not in _VALID_PADDING_MODES: + raise ValueError( + f"padding_mode must be one of {_VALID_PADDING_MODES}, got '{padding_mode}'" + ) + if isinstance(padding, str): + if padding not in _VALID_PADDING_STRINGS: + raise ValueError( + f"padding must be one of {_VALID_PADDING_STRINGS}, got '{padding}'" + ) + if padding == "same" and any(s != 1 for s in stride): + raise ValueError("padding='same' is not supported for strided convolutions") + + +class Conv2dLayer(nn.Module): + """Drop-in replacement for nn.Conv2d. Linear optimization disabled by default.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int, int], + stride: int | tuple[int, int] = 1, + padding: int | tuple[int, int] | str = 0, + dilation: int | tuple[int, int] = 1, + groups: int = 1, + bias: bool = True, + padding_mode: str = "zeros", + disable_linear: bool = True, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.kernel_size = _tuplify(kernel_size, 2) + self.stride = _tuplify(stride, 2) + self.dilation = _tuplify(dilation, 2) + self.groups = groups + self.padding_mode = padding_mode + + _validate_conv_args( + in_channels, out_channels, groups, padding, padding_mode, self.stride + ) + + if isinstance(padding, str): + self.padding = (0, 0) if padding == "valid" else padding + else: + self.padding = _tuplify(padding, 2) + + # Pre-compute pad tuple for padding_mode != "zeros" (mirrors nn.Conv2d). + # When padding="same", we need numeric values for F.pad; + # when padding is already numeric, _reverse_repeat_tuple handles it. + if isinstance(self.padding, str): + self._reversed_padding_repeated_twice = _compute_same_padding_for_pad( + self.kernel_size, self.dilation + ) + else: + self._reversed_padding_repeated_twice = _reverse_repeat_tuple(self.padding) + + padding_tuple = self.padding if isinstance(self.padding, tuple) else (1, 1) + self.enable_linear = not disable_linear and _check_enable_linear( + self.kernel_size, self.stride, padding_tuple, self.dilation, groups + ) + + self.weight = nn.Parameter( + torch.empty(out_channels, in_channels // groups, *self.kernel_size) + ) + if bias: + self.bias = nn.Parameter(torch.empty(out_channels)) + else: + self.register_parameter("bias", None) + + self._reset_parameters() + + def _reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + if self.bias is not None: + fan_in = nn.init._calculate_correct_fan(self.weight, "fan_in") + bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 + nn.init.uniform_(self.bias, -bound, bound) + + def _forward_mulmat(self, x: torch.Tensor) -> torch.Tensor: + K1, K2 = self.kernel_size + x = x.unfold(2, K1, K1).unfold(3, K2, K2) + N, _, Hp, Wp = x.shape[:4] + x = x.permute(0, 2, 3, 1, 4, 5).reshape(N, Hp, Wp, -1) + x = F.linear(x, self.weight.reshape(self.out_channels, -1), self.bias) + return x.permute(0, 3, 1, 2) + + def _forward_conv(self, x: torch.Tensor) -> torch.Tensor: + if self.padding_mode != "zeros": + return F.conv2d( + F.pad(x, self._reversed_padding_repeated_twice, mode=self.padding_mode), + self.weight, + self.bias, + self.stride, + (0, 0), + self.dilation, + self.groups, + ) + return F.conv2d( + x, + self.weight, + self.bias, + self.stride, + self.padding, + self.dilation, + self.groups, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.enable_linear: + return self._forward_mulmat(x) + return self._forward_conv(x) + + +class Conv3dLayer(nn.Module): + """Drop-in replacement for nn.Conv3d with automatic linear optimization.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int | tuple[int, int, int], + stride: int | tuple[int, int, int] = 1, + padding: int | tuple[int, int, int] | str = 0, + dilation: int | tuple[int, int, int] = 1, + groups: int = 1, + bias: bool = True, + padding_mode: str = "zeros", + disable_linear: bool = False, + ): + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.kernel_size = _tuplify(kernel_size, 3) + self.stride = _tuplify(stride, 3) + self.dilation = _tuplify(dilation, 3) + self.groups = groups + self.padding_mode = padding_mode + + _validate_conv_args( + in_channels, out_channels, groups, padding, padding_mode, self.stride + ) + + if isinstance(padding, str): + self.padding = (0, 0, 0) if padding == "valid" else padding + else: + self.padding = _tuplify(padding, 3) + + if isinstance(self.padding, str): + self._reversed_padding_repeated_twice = _compute_same_padding_for_pad( + self.kernel_size, self.dilation + ) + else: + self._reversed_padding_repeated_twice = _reverse_repeat_tuple(self.padding) + + padding_tuple = self.padding if isinstance(self.padding, tuple) else (1, 1, 1) + self.enable_linear = not disable_linear and _check_enable_linear( + self.kernel_size, self.stride, padding_tuple, self.dilation, groups + ) + + self.weight = nn.Parameter( + torch.empty(out_channels, in_channels // groups, *self.kernel_size) + ) + if bias: + self.bias = nn.Parameter(torch.empty(out_channels)) + else: + self.register_parameter("bias", None) + + self._reset_parameters() + + def _reset_parameters(self): + nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + if self.bias is not None: + fan_in = nn.init._calculate_correct_fan(self.weight, "fan_in") + bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 + nn.init.uniform_(self.bias, -bound, bound) + + def _forward_mulmat(self, x: torch.Tensor) -> torch.Tensor: + K1, K2, K3 = self.kernel_size + x = x.unfold(2, K1, K1).unfold(3, K2, K2).unfold(4, K3, K3) + N, Dp, Hp, Wp = x.shape[0], x.shape[2], x.shape[3], x.shape[4] + x = x.permute(0, 2, 3, 4, 1, 5, 6, 7).reshape(N, Dp, Hp, Wp, -1) + x = F.linear(x, self.weight.reshape(self.out_channels, -1), self.bias) + return x.permute(0, 4, 1, 2, 3) + + def _forward_conv(self, x: torch.Tensor) -> torch.Tensor: + if self.padding_mode != "zeros": + return F.conv3d( + F.pad(x, self._reversed_padding_repeated_twice, mode=self.padding_mode), + self.weight, + self.bias, + self.stride, + (0, 0, 0), + self.dilation, + self.groups, + ) + return F.conv3d( + x, + self.weight, + self.bias, + self.stride, + self.padding, + self.dilation, + self.groups, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.enable_linear: + return self._forward_mulmat(x) + return self._forward_conv(x) diff --git a/python/tokenspeed/runtime/layers/deepseek_v4_mhc.py b/python/tokenspeed/runtime/layers/deepseek_v4_mhc.py new file mode 100644 index 0000000..f2c3f85 --- /dev/null +++ b/python/tokenspeed/runtime/layers/deepseek_v4_mhc.py @@ -0,0 +1,496 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# Portions copyright the vLLM project contributors under Apache-2.0. + +from __future__ import annotations + +from functools import cache + +import torch +import triton +import triton.language as tl + +from tokenspeed.runtime.utils import ceil_div + +try: + from tokenspeed_kernel.thirdparty import deep_gemm +except Exception: + deep_gemm = None # type: ignore[assignment] + + +@cache +def _compute_num_split(block_k: int, k: int | None, grid_size: int) -> int: + device_props = torch.cuda.get_device_properties(0) + split_k = device_props.multi_processor_count // grid_size + if k is not None: + num_block_k = ceil_div(k, block_k) + split_k = min(split_k, num_block_k // 4) + return max(split_k, 1) + + +@triton.jit +def _load_reduced_mix( + gemm_out_mul, + token_id, + mix_id: tl.constexpr, + num_tokens, + hc_mult3: tl.constexpr, + n_splits: tl.constexpr, +): + value = tl.full((), 0.0, tl.float32) + for split_id in tl.static_range(0, n_splits): + offset = split_id * num_tokens * hc_mult3 + token_id * hc_mult3 + mix_id + value += tl.load(gemm_out_mul + offset) + return value + + +@triton.jit +def _mhc_pre_mix_triton_kernel( + gemm_out_mul, + gemm_out_sqrsum, + hc_scale, + hc_base, + pre_mix, + post_mix, + comb_mix, + hidden_size: tl.constexpr, + rms_eps: tl.constexpr, + hc_eps: tl.constexpr, + sinkhorn_iters: tl.constexpr, + n_splits: tl.constexpr, + hc_mult: tl.constexpr, + hc_mult2: tl.constexpr, + hc_mult3: tl.constexpr, + block_comb: tl.constexpr, + num_tokens, +): + token_id = tl.program_id(0) + + rms_sum = tl.full((), 0.0, tl.float32) + for split_id in tl.static_range(0, n_splits): + rms_sum += tl.load(gemm_out_sqrsum + split_id * num_tokens + token_id) + rms = tl.rsqrt(rms_sum / (hc_mult * hidden_size) + rms_eps) + + pre_scale = tl.load(hc_scale) + for hc_id in tl.static_range(0, hc_mult): + mix = _load_reduced_mix( + gemm_out_mul, + token_id, + hc_id, + num_tokens, + hc_mult3, + n_splits, + ) + pre = tl.sigmoid(mix * rms * pre_scale + tl.load(hc_base + hc_id)) + hc_eps + tl.store(pre_mix + token_id * hc_mult + hc_id, pre) + + post_scale = tl.load(hc_scale + 1) + for hc_id in tl.static_range(0, hc_mult): + mix = _load_reduced_mix( + gemm_out_mul, + token_id, + hc_mult + hc_id, + num_tokens, + hc_mult3, + n_splits, + ) + post = ( + tl.sigmoid(mix * rms * post_scale + tl.load(hc_base + hc_mult + hc_id)) + * 2.0 + ) + tl.store(post_mix + token_id * hc_mult + hc_id, post) + + comb_offsets = tl.arange(0, block_comb) + comb_mask = comb_offsets < hc_mult2 + comb_scale = tl.load(hc_scale + 2) + comb_mix_values = tl.zeros((block_comb,), tl.float32) + for split_id in tl.static_range(0, n_splits): + split_base = split_id * num_tokens * hc_mult3 + token_id * hc_mult3 + comb_mix_values += tl.load( + gemm_out_mul + split_base + hc_mult * 2 + comb_offsets, + mask=comb_mask, + other=0.0, + ) + comb_values = comb_mix_values * rms * comb_scale + tl.load( + hc_base + hc_mult * 2 + comb_offsets, mask=comb_mask, other=0.0 + ) + rows = comb_offsets // hc_mult + cols = comb_offsets - rows * hc_mult + active = comb_mask + + for row_id in tl.static_range(0, hc_mult): + row_values = tl.where((rows == row_id) & active, comb_values, -float("inf")) + row_max = tl.max(row_values, axis=0) + comb_values = tl.where( + (rows == row_id) & active, tl.exp(comb_values - row_max), comb_values + ) + for row_id in tl.static_range(0, hc_mult): + row_sum = tl.sum(tl.where((rows == row_id) & active, comb_values, 0.0), axis=0) + comb_values = tl.where( + (rows == row_id) & active, comb_values / row_sum + hc_eps, comb_values + ) + for col_id in tl.static_range(0, hc_mult): + col_sum = tl.sum(tl.where((cols == col_id) & active, comb_values, 0.0), axis=0) + comb_values = tl.where( + (cols == col_id) & active, + comb_values / (col_sum + hc_eps), + comb_values, + ) + + for _ in tl.static_range(1, sinkhorn_iters): + for row_id in tl.static_range(0, hc_mult): + row_sum = tl.sum( + tl.where((rows == row_id) & active, comb_values, 0.0), axis=0 + ) + comb_values = tl.where( + (rows == row_id) & active, + comb_values / (row_sum + hc_eps), + comb_values, + ) + for col_id in tl.static_range(0, hc_mult): + col_sum = tl.sum( + tl.where((cols == col_id) & active, comb_values, 0.0), axis=0 + ) + comb_values = tl.where( + (cols == col_id) & active, + comb_values / (col_sum + hc_eps), + comb_values, + ) + + tl.store( + comb_mix + token_id * hc_mult2 + comb_offsets, + comb_values, + mask=comb_mask, + ) + + +@triton.jit +def _mhc_pre_layer_triton_kernel( + pre_mix, + residual, + layer_input, + hidden_size: tl.constexpr, + hc_mult: tl.constexpr, + block_h: tl.constexpr, +): + token_id = tl.program_id(0) + hidden_block_id = tl.program_id(1) + + hidden_offsets = hidden_block_id * block_h + tl.arange(0, block_h) + hidden_mask = hidden_offsets < hidden_size + layer_acc = tl.zeros((block_h,), tl.float32) + for hc_id in tl.static_range(0, hc_mult): + pre = tl.load(pre_mix + token_id * hc_mult + hc_id).to(tl.float32) + residual_offsets = ( + token_id * hc_mult * hidden_size + hc_id * hidden_size + hidden_offsets + ) + residual_values = tl.load( + residual + residual_offsets, mask=hidden_mask, other=0.0 + ).to(tl.float32) + layer_acc += pre * residual_values + tl.store( + layer_input + token_id * hidden_size + hidden_offsets, + layer_acc, + mask=hidden_mask, + ) + + +@triton.jit +def _mhc_post_triton_kernel( + comb, + residual, + post, + hidden_states, + out, + hidden_size: tl.constexpr, + hc_mult: tl.constexpr, + block_h: tl.constexpr, +): + token_id = tl.program_id(0) + hidden_block_id = tl.program_id(1) + hidden_offsets = hidden_block_id * block_h + tl.arange(0, block_h) + hidden_mask = hidden_offsets < hidden_size + hidden_values = tl.load( + hidden_states + token_id * hidden_size + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + + for out_hc in tl.static_range(0, hc_mult): + acc = tl.load(post + token_id * hc_mult + out_hc).to(tl.float32) * hidden_values + for in_hc in tl.static_range(0, hc_mult): + comb_value = tl.load( + comb + token_id * hc_mult * hc_mult + in_hc * hc_mult + out_hc + ).to(tl.float32) + residual_values = tl.load( + residual + + token_id * hc_mult * hidden_size + + in_hc * hidden_size + + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + acc += comb_value * residual_values + tl.store( + out + + token_id * hc_mult * hidden_size + + out_hc * hidden_size + + hidden_offsets, + acc, + mask=hidden_mask, + ) + + +@triton.jit +def _mhc_post_hc4_triton_kernel( + comb, + residual, + post, + hidden_states, + out, + hidden_size: tl.constexpr, + block_h: tl.constexpr, +): + token_id = tl.program_id(0) + hidden_block_id = tl.program_id(1) + hidden_offsets = hidden_block_id * block_h + tl.arange(0, block_h) + hidden_mask = hidden_offsets < hidden_size + token_hidden_offset = token_id * hidden_size + token_residual_offset = token_id * 4 * hidden_size + + hidden_values = tl.load( + hidden_states + token_hidden_offset + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + + post_base = token_id * 4 + acc0 = tl.load(post + post_base + 0).to(tl.float32) * hidden_values + acc1 = tl.load(post + post_base + 1).to(tl.float32) * hidden_values + acc2 = tl.load(post + post_base + 2).to(tl.float32) * hidden_values + acc3 = tl.load(post + post_base + 3).to(tl.float32) * hidden_values + + comb_base = token_id * 16 + for in_hc in tl.static_range(0, 4): + residual_values = tl.load( + residual + token_residual_offset + in_hc * hidden_size + hidden_offsets, + mask=hidden_mask, + other=0.0, + ).to(tl.float32) + comb_row = comb_base + in_hc * 4 + acc0 += tl.load(comb + comb_row + 0).to(tl.float32) * residual_values + acc1 += tl.load(comb + comb_row + 1).to(tl.float32) * residual_values + acc2 += tl.load(comb + comb_row + 2).to(tl.float32) * residual_values + acc3 += tl.load(comb + comb_row + 3).to(tl.float32) * residual_values + + tl.store( + out + token_residual_offset + hidden_offsets, + acc0, + mask=hidden_mask, + ) + tl.store( + out + token_residual_offset + hidden_size + hidden_offsets, + acc1, + mask=hidden_mask, + ) + tl.store( + out + token_residual_offset + hidden_size * 2 + hidden_offsets, + acc2, + mask=hidden_mask, + ) + tl.store( + out + token_residual_offset + hidden_size * 3 + hidden_offsets, + acc3, + mask=hidden_mask, + ) + + +def mhc_fused_hc( + x_prev: torch.Tensor, + residual_prev: torch.Tensor, + post_prev: torch.Tensor, + comb_prev: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_eps: float, + sinkhorn_iters: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused post_mapping(prev) + pre_mapping(curr). + + Returns (residual_cur, layer_input, post_cur, comb_cur). + """ + residual_cur = mhc_post(x_prev, residual_prev, post_prev, comb_prev) + layer_input, post_cur, comb_cur = mhc_pre( + residual_cur, + fn, + hc_scale, + hc_base, + rms_eps, + hc_eps, + sinkhorn_iters, + ) + return residual_cur, layer_input, post_cur, comb_cur + + +def mhc_pre( + residual: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_eps: float, + sinkhorn_iters: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if residual.dtype != torch.bfloat16 or fn.dtype != torch.float32: + raise RuntimeError("fast mHC requires bf16 residual and fp32 weights") + if not residual.is_cuda: + raise RuntimeError("fast mHC requires CUDA tensors") + + if deep_gemm is None: + raise RuntimeError("deep_gemm.tf32_hc_prenorm_gemm is unavailable") + + hc_mult = residual.shape[-2] + hidden_size = residual.shape[-1] + hc_mult2 = hc_mult * hc_mult + hc_mult3 = hc_mult * 2 + hc_mult2 + hc_hidden_size = hc_mult * hidden_size + outer_shape = residual.shape[:-2] + residual_flat = residual.view(-1, hc_mult, hidden_size) + num_tokens = residual_flat.shape[0] + if num_tokens == 0: + return ( + residual.new_empty(*outer_shape, hidden_size), + torch.empty( + *outer_shape, + hc_mult, + 1, + dtype=torch.float32, + device=residual.device, + ), + torch.empty( + *outer_shape, + hc_mult, + hc_mult, + dtype=torch.float32, + device=residual.device, + ), + ) + + block_k = 64 + block_m = 64 + n_splits = _compute_num_split( + block_k, hc_hidden_size, ceil_div(num_tokens, block_m) + ) + + post_mix = torch.empty( + num_tokens, hc_mult, dtype=torch.float32, device=residual.device + ) + pre_mix = torch.empty( + num_tokens, hc_mult, dtype=torch.float32, device=residual.device + ) + comb_mix = torch.empty( + num_tokens, hc_mult2, dtype=torch.float32, device=residual.device + ) + layer_input = torch.empty( + num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device + ) + gemm_out_mul = torch.empty( + n_splits, num_tokens, hc_mult3, dtype=torch.float32, device=residual.device + ) + gemm_out_sqrsum = torch.empty( + n_splits, num_tokens, dtype=torch.float32, device=residual.device + ) + + deep_gemm.tf32_hc_prenorm_gemm( + residual_flat.view(num_tokens, hc_hidden_size), + fn, + gemm_out_mul, + gemm_out_sqrsum, + n_splits, + ) + block_h = 1024 + block_comb = triton.next_power_of_2(hc_mult2) + _mhc_pre_mix_triton_kernel[(num_tokens,)]( + gemm_out_mul, + gemm_out_sqrsum, + hc_scale, + hc_base, + pre_mix, + post_mix, + comb_mix, + hidden_size=hidden_size, + rms_eps=rms_eps, + hc_eps=hc_eps, + sinkhorn_iters=sinkhorn_iters, + n_splits=n_splits, + hc_mult=hc_mult, + hc_mult2=hc_mult2, + hc_mult3=hc_mult3, + block_comb=block_comb, + num_tokens=num_tokens, + num_warps=1, + ) + _mhc_pre_layer_triton_kernel[(num_tokens, triton.cdiv(hidden_size, block_h))]( + pre_mix, + residual_flat, + layer_input, + hidden_size=hidden_size, + hc_mult=hc_mult, + block_h=block_h, + num_warps=4, + ) + + return ( + layer_input.view(*outer_shape, hidden_size), + post_mix.view(*outer_shape, hc_mult, 1), + comb_mix.view(*outer_shape, hc_mult, hc_mult), + ) + + +def mhc_post( + hidden_states: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, +) -> torch.Tensor: + if not hidden_states.is_cuda: + raise RuntimeError("fast mHC requires CUDA tensors") + if residual.numel() == 0: + return torch.empty_like(residual) + out = torch.empty_like(residual) + hc_mult = residual.shape[-2] + hidden_size = residual.shape[-1] + residual_flat = residual.view(-1, hc_mult, hidden_size) + hidden_states_flat = hidden_states.view(-1, hidden_size) + post_flat = post.view(-1, hc_mult) + comb_flat = comb.view(-1, hc_mult, hc_mult) + num_tokens = residual_flat.shape[0] + if hc_mult == 4: + block_h = 256 + _mhc_post_hc4_triton_kernel[(num_tokens, triton.cdiv(hidden_size, block_h))]( + comb_flat, + residual_flat, + post_flat, + hidden_states_flat, + out, + hidden_size=hidden_size, + block_h=block_h, + num_warps=4, + ) + return out + + block_h = 1024 + _mhc_post_triton_kernel[(num_tokens, triton.cdiv(hidden_size, block_h))]( + comb_flat, + residual_flat, + post_flat, + hidden_states_flat, + out, + hidden_size=hidden_size, + hc_mult=hc_mult, + block_h=block_h, + num_warps=4, + ) + return out diff --git a/python/tokenspeed/runtime/layers/dense/__init__.py b/python/tokenspeed/runtime/layers/dense/__init__.py new file mode 100644 index 0000000..fea3eee --- /dev/null +++ b/python/tokenspeed/runtime/layers/dense/__init__.py @@ -0,0 +1,32 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. +from tokenspeed.runtime.layers.dense.fp8 import Fp8LinearMethod +from tokenspeed.runtime.layers.dense.mxfp4 import Mxfp4LinearMethod +from tokenspeed.runtime.layers.dense.nvfp4 import Nvfp4LinearMethod +from tokenspeed.runtime.layers.dense.unquant import UnquantizedLinearMethod +from tokenspeed.runtime.layers.dense.w8a8_fp8 import W8A8Fp8LinearMethod + +__all__ = [ + "Fp8LinearMethod", + "Mxfp4LinearMethod", + "Nvfp4LinearMethod", + "UnquantizedLinearMethod", + "W8A8Fp8LinearMethod", +] diff --git a/python/tokenspeed/runtime/layers/dense/fp8.py b/python/tokenspeed/runtime/layers/dense/fp8.py new file mode 100644 index 0000000..f5fc257 --- /dev/null +++ b/python/tokenspeed/runtime/layers/dense/fp8.py @@ -0,0 +1,377 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + + +import logging + +import tokenspeed_kernel +import torch +from tokenspeed_kernel.ops.gemm.fp8_utils import ( + per_token_group_quant_fp8, + per_token_quant_fp8, + static_quant_fp8, +) +from tokenspeed_kernel.platform import Platform +from torch.nn.parameter import Parameter + +logger = logging.getLogger(__name__) + +try: + from tokenspeed_kernel.thirdparty.deep_gemm import ceil_to_ue8m0 as _ceil_to_ue8m0 + from tokenspeed_kernel.thirdparty.deep_gemm import ( + transform_sf_into_required_layout as _transform_sf, + ) +except ImportError: + _ceil_to_ue8m0 = None + _transform_sf = None + +from tokenspeed.runtime.layers.dense.utils import normalize_e4m3fn_to_e4m3fnuz +from tokenspeed.runtime.layers.parameter import ( + BlockQuantScaleParameter, + ModelWeightParameter, + PerTensorScaleParameter, +) +from tokenspeed.runtime.layers.quantization.base_config import LinearMethodBase +from tokenspeed.runtime.layers.quantization.fp8 import Fp8Config +from tokenspeed.runtime.layers.quantization.utils import convert_to_channelwise + +platform = Platform.get() + + +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: Fp8Config): + self.quant_config = quant_config + self.block_quant = self.quant_config.weight_block_size is not 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, + ): + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + + 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 input_size > input_size_per_partition: + 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 ( + output_size > output_size_per_partition + 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"): + if self.quant_config.activation_scheme != "dynamic": + raise ValueError( + "Block FP8 requires dynamic activation quantization." + ) + elif hasattr(self.quant_config, "linear_activation_scheme"): + if self.quant_config.linear_activation_scheme != "dynamic": + raise ValueError( + "Block FP8 requires dynamic linear activation quantization." + ) + 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: torch.nn.Module) -> None: + if self.block_quant: + # If ROCm, normalize the weights and scales to e4m3fnuz + if platform.is_fp8e4m3fnuz: + # 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 + else: + weight, weight_scale = layer.weight.data, layer.weight_scale_inv.data + layer.weight.data = weight.data + layer.weight_scale_inv.data = weight_scale.data + layer._use_deep_gemm_fp8 = False + is_bmm = getattr(layer, "is_bmm", False) + is_ue8m0 = getattr(self.quant_config, "scale_fmt", None) == "ue8m0" + if _transform_sf is not None and _ceil_to_ue8m0 is not None and is_ue8m0: + N, K = layer.weight.shape + block_n, block_k = self.quant_config.weight_block_size + if is_bmm: + # Grouped (batched) projection (V4 attention wo_a, weight + # [groups * n, K], consumed per group as [n, K]). Transform + # the block scale into the deep_gemm MN-major layout with the + # group axis so deep_gemm.fp8_einsum("bhr,hdr->bhd") runs the + # output projection as one native FP8 GEMM (no FP32 dequant). + # recipe is (1, block_n, block_k) at load; the runtime einsum + # uses (1, 1, block_n) on SM100. + g = layer.bmm_batch_size + n = N // g + if n % block_n == 0 and K % block_k == 0: + sf = _ceil_to_ue8m0(layer.weight_scale_inv.data).view( + g, n // block_n, K // block_k + ) + layer.weight_scale_inv.data = _transform_sf( + sf=sf, + mn=n, + k=K, + recipe=(1, block_n, block_k), + num_groups=g, + is_sfa=False, + ) + layer._deep_gemm_block_size = [block_n, block_k] + layer._use_deep_gemm_fp8 = True + elif N % 64 == 0 and K % 128 == 0: + sf = _ceil_to_ue8m0(layer.weight_scale_inv.data) + layer.weight_scale_inv.data = _transform_sf( + sf=sf, + mn=N, + k=K, + recipe=(1, block_n, block_k), + is_sfa=False, + ) + layer._use_deep_gemm_fp8 = True + if is_bmm and not layer._use_deep_gemm_fp8: + # The is_bmm runtime path (DeepSeek-V4 o_proj) has no FP32 + # fallback, so fail fast at load with a clear message instead of + # a cryptic AttributeError on the first forward. + raise RuntimeError( + "is_bmm weight requires the deep_gemm FP8 block-scale path " + "but it could not be prepared (deep_gemm_available=" + f"{_transform_sf is not None}, ue8m0={is_ue8m0}, " + f"weight={tuple(layer.weight.shape)}); ensure FP8 block-quant " + "ue8m0 weights with block-aligned dims and deep_gemm installed." + ) + 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: + # apply per-channel quantization default as + qweight, weight_scale = per_token_group_quant_fp8( + layer.weight, layer.weight.shape[-1] + ) + weight_scale = weight_scale.t().contiguous() + + # 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 + ) + + weight = layer.weight + weight_scale = convert_to_channelwise( + layer.weight_scale, 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 + ) + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + block_scale: torch.Tensor | None = None, + output_dtype: torch.dtype | None = None, + ) -> torch.Tensor: + + if self.block_quant: + input_2d = x.view(-1, x.shape[-1]) + output_shape = [*x.shape[:-1], layer.weight.shape[0]] + output_dtype = output_dtype or x.dtype + + override = ( + "deep_gemm_mm_fp8_blockscale" + if getattr(layer, "_use_deep_gemm_fp8", False) + else None + ) + output = tokenspeed_kernel.mm( + input_2d, + layer.weight, + A_scales=block_scale, + B_scales=layer.weight_scale_inv, + bias=bias, + out_dtype=output_dtype, + quant="mxfp8", + block_size=self.quant_config.weight_block_size, + override=override, + ) + return output.to(dtype=output_dtype).view(*output_shape) + else: + input = x + weight = layer.weight + weight_scale = layer.weight_scale + input_scale = layer.input_scale + + # 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 input_scale is not None: + if input_scale.numel() != 1: + raise ValueError( + f"input_scale must contain exactly one value, got {input_scale.numel()}." + ) + qinput, x_scale = static_quant_fp8(input_2d, input_scale) + else: + qinput, x_scale = per_token_quant_fp8(input_2d) + + qinput = qinput.view(-1, qinput.shape[-1]) + + output = tokenspeed_kernel.mm( + qinput, + weight, + A_scales=x_scale, + B_scales=weight_scale, + out_dtype=input.dtype, + ) + if bias is not None: + output = output + bias + return output.view(*output_shape) diff --git a/python/tokenspeed/runtime/layers/dense/mxfp4.py b/python/tokenspeed/runtime/layers/dense/mxfp4.py new file mode 100644 index 0000000..6b5a79c --- /dev/null +++ b/python/tokenspeed/runtime/layers/dense/mxfp4.py @@ -0,0 +1,152 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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 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. + +"""MXFP4 dense linear bridge for checkpoint-serialized Kimi MLP weights.""" + +from __future__ import annotations + +import tokenspeed_kernel +import torch +from tokenspeed_kernel.ops.quantization.triton import mxfp4_quantize +from torch.nn.parameter import Parameter + +from tokenspeed.runtime.layers.quantization.base_config import QuantizeMethodBase + +MXFP4_BLOCK = 32 + + +class Mxfp4LinearMethod(QuantizeMethodBase): + """Packed MXFP4 dense weights. + + Kimi-K2.5 MXFP4 stores dense layer-0 MLP and MoE shared-expert MLP tensors + in the same packed FP4/e8m0 format as routed experts. Runtime activations + are quantized to packed MXFP4 before the dense GEMM so checkpoint weights + can stay packed in VRAM. + """ + + def __init__(self, quant_config): + self.quant_config = quant_config + self.group_size = getattr(quant_config, "group_size", MXFP4_BLOCK) + + 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 + _validate_mxfp4_partition(input_size_per_partition) + output_size_per_partition = sum(output_partition_sizes) + weight_loader = extra_weight_attrs.get("weight_loader") + scale_loader = _wrap_e8m0_scale_loader(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 + layer.params_dtype = params_dtype + layer.orig_dtype = params_dtype + + weight = Parameter( + torch.empty( + output_size_per_partition, + input_size_per_partition // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ) + weight.output_dim = 0 + weight.input_dim = 1 + if weight_loader is not None: + weight.weight_loader = weight_loader + layer.register_parameter("weight", weight) + + weight_scale = Parameter( + torch.empty( + output_size_per_partition, + input_size_per_partition // self.group_size, + dtype=torch.uint8, + ), + requires_grad=False, + ) + weight_scale.output_dim = 0 + weight_scale.input_dim = 1 + if scale_loader is not None: + weight_scale.weight_loader = scale_loader + layer.register_parameter("weight_scale", weight_scale) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + if getattr(layer, "_mxfp4_dense_processed", False): + return + layer.weight_triton_tensor = layer.weight.data + layer.weight_scale_triton_tensor = layer.weight_scale.data + layer._mxfp4_dense_processed = True + + def apply(self, layer, x, bias=None): + if not getattr(layer, "_mxfp4_dense_processed", False): + self.process_weights_after_loading(layer) + input_2d = x.reshape(-1, x.shape[-1]) + output_shape = (*x.shape[:-1], layer.output_size_per_partition) + input_quant, input_scale = mxfp4_quantize(input_2d) + output = tokenspeed_kernel.mm( + input_quant, + layer.weight_triton_tensor, + A_scales=input_scale, + B_scales=layer.weight_scale_triton_tensor, + bias=bias, + out_dtype=x.dtype, + quant="mxfp4", + expected_kernel_name="triton_mm_mxfp4", + ) + return output.reshape(*output_shape) + + +def _validate_mxfp4_partition(input_size_per_partition: int) -> None: + if input_size_per_partition % 2 != 0: + raise ValueError( + f"MXFP4 input partition {input_size_per_partition} must be divisible by 2" + ) + if input_size_per_partition % MXFP4_BLOCK != 0: + raise ValueError( + f"MXFP4 input partition {input_size_per_partition} must be divisible by " + f"{MXFP4_BLOCK}" + ) + + +def _wrap_e8m0_scale_loader(weight_loader): + if weight_loader is None: + return None + + def scale_loader(param, loaded_weight, *args, **kwargs): + e8m0_dtype = getattr(torch, "float8_e8m0fnu", None) + if ( + e8m0_dtype is not None + and param.dtype == torch.uint8 + and loaded_weight.dtype == e8m0_dtype + ): + loaded_weight = loaded_weight.view(torch.uint8) + return weight_loader(param, loaded_weight, *args, **kwargs) + + return scale_loader + + +__all__ = [ + "Mxfp4LinearMethod", +] diff --git a/python/tokenspeed/runtime/layers/dense/nvfp4.py b/python/tokenspeed/runtime/layers/dense/nvfp4.py new file mode 100644 index 0000000..557e219 --- /dev/null +++ b/python/tokenspeed/runtime/layers/dense/nvfp4.py @@ -0,0 +1,274 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import logging + +import tokenspeed_kernel +import torch +from tokenspeed_kernel.ops.quantization.flashinfer import fp4_quantize +from torch.nn.parameter import Parameter + +from tokenspeed.runtime.layers.quantization.base_config import QuantizeMethodBase + +logger = logging.getLogger(__name__) + + +def _pdl_enabled() -> bool: + from tokenspeed.runtime.utils.pdl import pdl_enabled + + return pdl_enabled() + + +class Nvfp4LinearMethod(QuantizeMethodBase): + """Linear method for NVFP4 quantization. + + Weight structure: + - weight: uint8 [output_size, input_size // 2] (packed FP4) + - weight_scale: float8_e4m3fn [output_size, input_size // group_size] + - weight_scale_2: float32 scalar (per-tensor) + - input_scale: float32 scalar (per-tensor) + """ + + def __init__(self, quant_config): + self.quant_config = quant_config + self.group_size = 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, + ): + 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 + + # FP4 packed weight: 2 values per byte, input_dim halved + weight = Parameter( + torch.empty( + output_size_per_partition, + input_size_per_partition // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ) + # Set attributes for TP sharding in weight_loader + weight.output_dim = 0 + weight.input_dim = 1 + if weight_loader: + weight.weight_loader = weight_loader + layer.register_parameter("weight", weight) + + # Block scales: one per group_size elements + weight_scale = Parameter( + torch.empty( + output_size_per_partition, + input_size_per_partition // self.group_size, + dtype=torch.float8_e4m3fn, + ), + requires_grad=False, + ) + weight_scale.output_dim = 0 + weight_scale.input_dim = 1 + if weight_loader: + weight_scale.weight_loader = weight_loader + layer.register_parameter("weight_scale", weight_scale) + + # Per-tensor scales: scalar per partition, use needs_scalar_to_array for fused loading + input_scale = Parameter( + torch.full( + (len(output_partition_sizes),), + torch.finfo(torch.float32).min, + dtype=torch.float32, + ), + requires_grad=False, + ) + input_scale.needs_scalar_to_array = True + if weight_loader: + input_scale.weight_loader = weight_loader + layer.register_parameter("input_scale", input_scale) + + weight_scale_2 = Parameter( + torch.full( + (len(output_partition_sizes),), + torch.finfo(torch.float32).min, + dtype=torch.float32, + ), + requires_grad=False, + ) + weight_scale_2.needs_scalar_to_array = True + if weight_loader: + weight_scale_2.weight_loader = weight_loader + layer.register_parameter("weight_scale_2", weight_scale_2) + + def process_weights_after_loading(self, layer): + """Compute alpha and input_scale_inv, swizzle block scales.""" + logger.debug( + "[FP4_DENSE_POSTLOAD] w=%s(%s) ws=%s is=%s ws2=%s", + layer.weight.shape, + layer.weight.dtype, + layer.weight_scale.shape, + layer.input_scale, + layer.weight_scale_2, + ) + input_scale = layer.input_scale.max().to(torch.float32) + weight_scale_2 = layer.weight_scale_2.max().to(torch.float32) + layer.input_scale = Parameter(input_scale, requires_grad=False) + layer.weight_scale_2 = Parameter(weight_scale_2, requires_grad=False) + layer.alpha = Parameter(input_scale * weight_scale_2, requires_grad=False) + layer.input_scale_inv = Parameter( + (1.0 / input_scale).to(torch.float32), requires_grad=False + ) + if layer.interleave_linear_and_gate: + gate_weight, linear_weight = layer.weight.chunk(2, dim=0) + linear_gate_weight = torch.cat((linear_weight, gate_weight), dim=0) + layer.weight_swiglu_interleaved = Parameter( + interleave_linear_and_gate( + linear_gate_weight, + group_size=64, + dim=0, + ), + requires_grad=False, + ) + # layer.weight_scale is the canonical unswizzled [N, K/group] + # tensor. Reorder gate/linear first, then swizzle for the CUTE + # kernel; chunking layer.weight_scale_interleaved would be wrong. + gate_scale, linear_scale = layer.weight_scale.chunk(2, dim=0) + linear_gate_scale = torch.cat((linear_scale, gate_scale), dim=0) + layer.weight_scale_swiglu_interleaved = Parameter( + swizzle_blockscale_2d( + interleave_linear_and_gate( + linear_gate_scale, + group_size=64, + dim=0, + ) + ), + requires_grad=False, + ) + del layer.weight + del layer.weight_scale + else: + # Swizzle block scales for CUTLASS + layer.weight_scale_interleaved = Parameter( + swizzle_blockscale_2d(layer.weight_scale), + requires_grad=False, + ) + del layer.weight_scale + + def apply(self, layer, x, bias=None): + """Forward pass: quantize input to FP4, run FP4 GEMM. + + ``x`` may be either a bf16/fp16 activation tensor (normal path) or a + pre-quantized ``(x_fp4, x_scale)`` tuple. + """ + w_n = layer.output_size_per_partition + + if isinstance(x, tuple): + # Pre-quantized path: no fp4_quantize launch. Output dtype is bf16. + x_fp4, x_scale = x + output_dtype = torch.bfloat16 + + else: + x_fp4, x_scale = fp4_quantize( + x, layer.input_scale_inv, enable_pdl=_pdl_enabled() + ) + output_dtype = x.dtype + + kernel_override = layer.override_kernel_name + out = tokenspeed_kernel.mm( + x_fp4, + layer.weight.T, + A_scales=x_scale, + B_scales=layer.weight_scale_interleaved.T, + bias=bias, + out_dtype=output_dtype, + alpha=layer.alpha, + quant="nvfp4", + enable_pdl=_pdl_enabled(), + override=kernel_override, + expected_kernel_name=kernel_override or "cublaslt_mm_nvfp4", + ) + return out.view(x_fp4.size(0), w_n) + + +# ------------------------------------------------------------------------- +# Utilities for FP4 linear method +# ------------------------------------------------------------------------- + + +def swizzle_blockscale_2d(scales): + """Swizzle 2D FP8 block scales for CUTLASS.""" + M, K = scales.shape + + def round_up(x, m): + return (x + m - 1) // m * m + + M_padded = round_up(M, 128) + K_padded = round_up(K, 4) + padded = torch.zeros((M_padded, K_padded), dtype=scales.dtype, device=scales.device) + padded[:M, :K] = scales + rows, cols = padded.shape + padded = padded.reshape(rows // 128, 4, 32, cols // 4, 4) + padded = padded.permute((0, 3, 2, 1, 4)) + return padded.contiguous().reshape(M_padded, K_padded) + + +def interleave_linear_and_gate( + tensor: torch.Tensor, + group_size: int = 64, + dim: int = 0, +) -> torch.Tensor: + """Interleave ``[linear all][gate all]`` as ``[linear chunk][gate chunk]``. + + This matches the FC1 GEMM+SwiGLU preprocessing layout expected by + ``cute_dsl_nvfp4_dense_gemm_swiglu_blackwell``. + """ + if tensor.ndim == 0: + raise ValueError("expected a tensor with at least one dimension") + + dim = dim % tensor.ndim + sizes = tensor.size() + dim_size = sizes[dim] + if dim_size % (group_size * 2) != 0: + raise ValueError( + f"dimension {dim} size {dim_size} must be divisible by " + f"2 * group_size={2 * group_size}" + ) + + prev_sizes = sizes[:dim] + post_sizes = sizes[dim + 1 :] + return ( + tensor.reshape( + *prev_sizes, + 2, + dim_size // (group_size * 2), + group_size, + *post_sizes, + ) + .transpose(dim, dim + 1) + .reshape(*sizes) + .contiguous() + ) diff --git a/python/tokenspeed/runtime/layers/dense/unquant.py b/python/tokenspeed/runtime/layers/dense/unquant.py new file mode 100644 index 0000000..0dcba78 --- /dev/null +++ b/python/tokenspeed/runtime/layers/dense/unquant.py @@ -0,0 +1,69 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + + +import tokenspeed_kernel +import torch +from torch.nn.parameter import Parameter + +from tokenspeed.runtime.layers.quantization.base_config import LinearMethodBase +from tokenspeed.runtime.utils import set_weight_attrs + + +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: + return + + def apply( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + + return tokenspeed_kernel.mm( + x, + layer.weight, + bias=bias, + ) diff --git a/python/tokenspeed/runtime/layers/dense/utils.py b/python/tokenspeed/runtime/layers/dense/utils.py new file mode 100644 index 0000000..2f169ad --- /dev/null +++ b/python/tokenspeed/runtime/layers/dense/utils.py @@ -0,0 +1,46 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + + +import torch + + +def normalize_e4m3fn_to_e4m3fnuz( + weight: torch.Tensor, + weight_scale: torch.Tensor, + input_scale: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + 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 diff --git a/python/tokenspeed/runtime/layers/dense/w8a8_fp8.py b/python/tokenspeed/runtime/layers/dense/w8a8_fp8.py new file mode 100755 index 0000000..13c9316 --- /dev/null +++ b/python/tokenspeed/runtime/layers/dense/w8a8_fp8.py @@ -0,0 +1,141 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import tokenspeed_kernel +import torch +from tokenspeed_kernel.ops.gemm.fp8_utils import ( + per_token_group_quant_fp8, + per_token_quant_fp8, +) +from tokenspeed_kernel.platform import Platform +from torch.nn.parameter import Parameter + +from tokenspeed.runtime.layers.dense.utils import normalize_e4m3fn_to_e4m3fnuz +from tokenspeed.runtime.layers.parameter import ( + ChannelQuantScaleParameter, + ModelWeightParameter, +) +from tokenspeed.runtime.layers.quantization.base_config import LinearMethodBase +from tokenspeed.runtime.layers.quantization.w8a8_fp8 import W8A8Fp8Config + +platform = Platform.get() + + +class W8A8Fp8LinearMethod(LinearMethodBase): + + def __init__(self, quantization_config: W8A8Fp8Config): + 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 platform.is_fp8e4m3fnuz: + 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: + # use 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() + + # 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: torch.Tensor | None = None, + ): + input = x + weight = layer.weight + weight_scale = layer.weight_scale + + # View input as 2D matrix for fp8 methods + input_2d = input.view(-1, input.shape[-1]) + output_shape = [*input.shape[:-1], weight.shape[1]] + + qinput, x_scale = per_token_quant_fp8(input_2d) + + qinput = qinput.view(-1, qinput.shape[-1]) + + output = tokenspeed_kernel.mm( + qinput, + weight, + A_scales=x_scale, + B_scales=weight_scale, + out_dtype=input.dtype, + ) + if bias is not None: + output = output + bias + return output.view(*output_shape) diff --git a/python/tokenspeed/runtime/layers/layernorm.py b/python/tokenspeed/runtime/layers/layernorm.py new file mode 100755 index 0000000..da7c027 --- /dev/null +++ b/python/tokenspeed/runtime/layers/layernorm.py @@ -0,0 +1,530 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Fused operators for normalization layers.""" + +import torch +import torch.nn as nn +from tokenspeed_kernel.ops.communication.triton import ( + allreduce_residual_rmsnorm as triton_allreduce_residual_rmsnorm, +) +from tokenspeed_kernel.ops.communication.trtllm import ( + allgather_dual_rmsnorm, +) +from tokenspeed_kernel.ops.communication.trtllm import ( + allreduce_residual_rmsnorm as trtllm_allreduce_residual_rmsnorm, +) +from tokenspeed_kernel.ops.communication.trtllm import ( + reducescatter_residual_rmsnorm, +) +from tokenspeed_kernel.platform import current_platform + +from tokenspeed.runtime.distributed.process_group_manager import ( + process_group_manager as pg_manager, +) +from tokenspeed.runtime.utils import ( + get_colorful_logger, +) +from tokenspeed.runtime.utils.env import global_server_args_dict +from tokenspeed.runtime.utils.pdl import pdl_enabled + +_is_amd = current_platform().is_amd + +if _is_amd: + from tokenspeed_kernel.ops.layernorm.triton import rmsnorm as triton_rmsnorm + from tokenspeed_kernel.ops.layernorm.triton import ( + rmsnorm_fused_parallel as triton_rmsnorm_fused_parallel, + ) +else: + from tokenspeed_kernel.ops.layernorm.cuda import rmsnorm_fused_parallel + from tokenspeed_kernel.ops.layernorm.flashinfer import ( + fused_add_rmsnorm, + gemma_fused_add_rmsnorm, + gemma_rmsnorm, + layernorm, + rmsnorm, + ) + + +logger = get_colorful_logger(__name__) + + +def _get_process_group(group: tuple[int, ...]): + return pg_manager.get_process_group("nccl", group) + + +class LayerNorm(nn.Module): + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size, dtype=torch.float32)) + self.bias = nn.Parameter(torch.zeros(hidden_size, dtype=torch.float32)) + self.variance_epsilon = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # There might be no tokens here (e.g. idle/padded graph rows). + if x.shape[0] == 0: + return x + if current_platform().is_nvidia: + return layernorm(x, self.weight, self.bias, self.variance_epsilon) + return nn.functional.layer_norm( + x.float(), + (x.shape[-1],), + self.weight, + self.bias, + self.variance_epsilon, + ).to(x.dtype) + + +class RMSNorm(torch.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: torch.Tensor | None = None, + inplace: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + # There might be no tokens here + if x.shape[0] == 0: + if residual is not None: + return x, residual + else: + return x + + if _is_amd: + if residual is not None: + if inplace: + raise ValueError( + "fused add rmsnorm does not support inplace operation" + ) + return triton_rmsnorm( + x, + self.weight.data, + self.variance_epsilon, + residual=residual, + ) + return triton_rmsnorm( + x, + self.weight.data, + self.variance_epsilon, + out=x if inplace else None, + ) + else: + if residual is not None: + if inplace: + raise ValueError( + "fused_add_rmsnorm does not support inplace operation" + ) + fused_add_rmsnorm( + x, + residual, + self.weight.data, + self.variance_epsilon, + enable_pdl=pdl_enabled(), + ) + return x, residual + out = rmsnorm( + x, + self.weight.data, + self.variance_epsilon, + out=x if inplace else None, + enable_pdl=pdl_enabled(), + ) + return out + + def forward_with_allreduce_fusion( + self, + rank: int, + group: tuple[int, ...], + x: torch.Tensor, + residual: torch.Tensor | None = None, + fuse_block_quant_fp8: bool = False, + residual_reduce_scattered: bool = False, + max_sm_to_use: int | None = None, + trigger_completion_at_end: bool = False, + has_partial_norm_out: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """ + Forward method with allreduce fusion, prioritizing flashinfer fused operations + """ + + if residual is not None: + + if len(group) > 1: + if _is_amd: + allreduce_residual_rmsnorm = triton_allreduce_residual_rmsnorm + else: + if not current_platform().is_nvidia: + raise RuntimeError("Allreduce RMSNorm requires NVIDIA or AMD.") + allreduce_residual_rmsnorm = trtllm_allreduce_residual_rmsnorm + fused_result = allreduce_residual_rmsnorm( + input_tensor=x, + residual=residual, + weight=self.weight, + rank=rank, + group=_get_process_group(group), + eps=self.variance_epsilon, + max_token_num=global_server_args_dict["comm_fusion_max_num_tokens"], + block_quant_fp8=fuse_block_quant_fp8, + residual_reduce_scattered=residual_reduce_scattered, + max_sm_to_use=max_sm_to_use, + trigger_completion_at_end=trigger_completion_at_end, + has_partial_norm_out=has_partial_norm_out, + launch_with_pdl=pdl_enabled(), + ) + if fused_result[0] is not None: + return fused_result + + result = self.forward(x, residual) + if isinstance(result, tuple): + return result[0], result[1], None + return result, None, None + + def forward_with_reducescatter_fusion( + self, + rank: int, + group: tuple[int, ...], + x: torch.Tensor, + residual: torch.Tensor | None = None, + fuse_block_quant_fp8: bool = False, + add_in: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """ + Forward method with reducescatter fusion, prioritizing flashinfer fused operations + """ + + if residual is not None: + + if len(group) > 1: + fused_result = reducescatter_residual_rmsnorm( + input_tensor=x, + residual=residual, + weight=self.weight, + rank=rank, + group=_get_process_group(group), + eps=self.variance_epsilon, + max_token_num=global_server_args_dict["comm_fusion_max_num_tokens"], + use_oneshot=True, + block_quant_fp8=fuse_block_quant_fp8, + add_in=add_in, + launch_with_pdl=pdl_enabled(), + ) + if fused_result[0] is not None: + return fused_result + + result = self.forward(x, residual) + if isinstance(result, tuple): + return result[0], result[1], None + return result, None, None + + +class GemmaRMSNorm(torch.nn.Module): + 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 + self.register_buffer("gemma_weight", self.weight.data + 1.0, persistent=False) + # (Chen-0210) Gemma weight = standard_weight + 1. Precompute once. + self.weight.weight_loader = self._weight_loader + + def _weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + if param.size() != loaded_weight.size(): + raise ValueError( + f"Shape mismatch: {param.size()} != {loaded_weight.size()}." + ) + param.data.copy_(loaded_weight) + self.gemma_weight = param.data + 1.0 + + def forward( + self, + x: torch.Tensor, + residual: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + if x.shape[0] == 0: + if residual is not None: + return x, residual + else: + return x + + if _is_amd: + if x.shape[0] == 0: + if residual is not None: + return x, residual + else: + return x + 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) + else: + if residual is not None: + gemma_fused_add_rmsnorm( + x, + residual, + self.weight.data, + self.variance_epsilon, + enable_pdl=pdl_enabled(), + ) + return x, residual + out = gemma_rmsnorm( + x, + self.weight.data, + self.variance_epsilon, + enable_pdl=pdl_enabled(), + ) + return out + + def forward_with_allreduce_fusion( + self, + rank: int, + group: tuple[int, ...], + x: torch.Tensor, + residual: torch.Tensor | None = None, + fuse_block_quant_fp8: bool = False, + residual_reduce_scattered: bool = False, + max_sm_to_use: int | None = None, + trigger_completion_at_end: bool = False, + has_partial_norm_out: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]: + """ + Forward method with allreduce fusion for GemmaRMSNorm. + Uses gemma_weight (= weight + 1.0) as gamma so that the standard + fused kernel computes x * (1 + weight) matching GemmaRMSNorm semantics. + """ + + if residual is not None: + + if len(group) > 1: + if _is_amd: + allreduce_residual_rmsnorm = triton_allreduce_residual_rmsnorm + else: + if not current_platform().is_nvidia: + raise RuntimeError("Allreduce RMSNorm requires NVIDIA or AMD.") + allreduce_residual_rmsnorm = trtllm_allreduce_residual_rmsnorm + fused_result = allreduce_residual_rmsnorm( + input_tensor=x, + residual=residual, + weight=self.gemma_weight, + rank=rank, + group=_get_process_group(group), + eps=self.variance_epsilon, + max_token_num=global_server_args_dict["comm_fusion_max_num_tokens"], + block_quant_fp8=fuse_block_quant_fp8, + residual_reduce_scattered=residual_reduce_scattered, + max_sm_to_use=max_sm_to_use, + trigger_completion_at_end=trigger_completion_at_end, + has_partial_norm_out=has_partial_norm_out, + launch_with_pdl=pdl_enabled(), + ) + if fused_result[0] is not None: + return fused_result + + result = self.forward(x, residual) + if isinstance(result, tuple): + return result[0], result[1], None + return result, None, None + + def forward_with_reducescatter_fusion( + self, + rank: int, + group: tuple[int, ...], + x: torch.Tensor, + residual: torch.Tensor | None = None, + fuse_block_quant_fp8: bool = False, + add_in: torch.Tensor | None = None, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """ + Forward method with reducescatter fusion for GemmaRMSNorm. + Uses gemma_weight (= weight + 1.0) as gamma so that the standard + fused kernel computes x * (1 + weight) matching GemmaRMSNorm semantics. + """ + + if residual is not None: + + if len(group) > 1: + fused_result = reducescatter_residual_rmsnorm( + input_tensor=x, + residual=residual, + weight=self.gemma_weight, + rank=rank, + group=_get_process_group(group), + eps=self.variance_epsilon, + max_token_num=global_server_args_dict["comm_fusion_max_num_tokens"], + use_oneshot=True, + block_quant_fp8=fuse_block_quant_fp8, + add_in=add_in, + launch_with_pdl=pdl_enabled(), + ) + if fused_result[0] is not None: + return fused_result + + result = self.forward(x, residual) + if isinstance(result, tuple): + return result[0], result[1], None + return result, None, None + + +class FusedRMSNorm(nn.Module): + """Fused RMSNorm layer for normalizing two tensors simultaneously. + + This layer wraps two independent RMSNorm layers (q_a and kv_a) and performs + fused normalization during forward pass. The RMSNorm layers are passed in as + parameters, allowing reuse of existing normalization layers. + """ + + def __init__( + self, + q_a_norm: RMSNorm, + kv_a_norm: RMSNorm, + ) -> None: + super().__init__() + self.q_a_norm = q_a_norm + self.kv_a_norm = kv_a_norm + + @property + def weight_q_a(self) -> nn.Parameter: + """Expose weight_q_a from q_a_norm for backward compatibility.""" + return self.q_a_norm.weight + + @property + def weight_kv_a(self) -> nn.Parameter: + """Expose weight_kv_a from kv_a_norm for backward compatibility.""" + return self.kv_a_norm.weight + + def forward( + self, + input_q_a: torch.Tensor, + input_kv_a: torch.Tensor, + output_q_a: torch.Tensor | None = None, + output_kv_a: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Normalize two tensors in parallel using fused computation. + + Args: + input_q_a: Q tensor to normalize + input_kv_a: KV tensor to normalize + + Returns: + Tuple of (normalized_q_a, normalized_kv_a) + """ + if _is_amd: + triton_rmsnorm_fused_parallel( + input1=input_q_a, + weight1=self.weight_q_a, + output1=output_q_a if output_q_a is not None else input_q_a, + input2=input_kv_a, + weight2=self.weight_kv_a, + output2=output_kv_a if output_kv_a is not None else input_kv_a, + eps=self.q_a_norm.variance_epsilon, + enable_pdl=pdl_enabled(), + ) + else: + rmsnorm_fused_parallel( + input1=input_q_a, + weight1=self.weight_q_a, + output1=output_q_a if output_q_a is not None else input_q_a, + input2=input_kv_a, + weight2=self.weight_kv_a, + output2=output_kv_a if output_kv_a is not None else input_kv_a, + eps=self.q_a_norm.variance_epsilon, + enable_pdl=pdl_enabled(), + ) + return input_q_a, input_kv_a + + def forward_with_allgather_fusion( + self, + rank: int, + group: tuple[int, ...], + qkv: torch.Tensor, + total_num_tokens: int, + fuse_block_quant_fp8: bool = False, + trigger_completion_at_end: bool = False, + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: + """ + Forward method with allgather fusion, performing allgather + dual RMSNorm + optional FP8 block quantization. + + This method fuses allgather communication with dual RMSNorm computation + and optional FP8 block-wise quantization in a single kernel launch. + + Args: + qkv: Input tensor to allgather, shape [num_token_current_rank, q_lora_rank + kv_lora_rank + qk_rope_head_dim] + fuse_block_quant_fp8: Whether to perform FP8 block-wise quantization on the first norm output + trigger_completion_at_end: Whether to trigger completion event at the end of kernel + + Returns: + Tuple of (allgather_out, quant_out, k_nope, block_scale): + - allgather_out: Gathered tensor, shape [num_token_all_group, hidden_dim] + - quant_out: FP8 quantized first norm output (q_contiguous), None if fuse_block_quant_fp8=False + - k_nope: Second norm output + - block_scale: Quantization scales, None if fuse_block_quant_fp8=False + """ + + if len(group) > 1: + fused_result = allgather_dual_rmsnorm( + qkv=qkv, + total_num_tokens=total_num_tokens, + rank=rank, + group=_get_process_group(group), + weight_q_a=self.weight_q_a, + weight_kv_a=self.weight_kv_a, + eps_q=self.q_a_norm.variance_epsilon, + eps_kv=self.kv_a_norm.variance_epsilon, + max_token_num=global_server_args_dict["comm_fusion_max_num_tokens"], + block_quant_fp8=fuse_block_quant_fp8, + trigger_completion_at_end=trigger_completion_at_end, + fp32_acc=False, + launch_with_pdl=pdl_enabled(), + ) + if fused_result[0] is not None: + return fused_result + + q_lora_rank = self.weight_q_a.shape[0] + kv_lora_rank = self.weight_kv_a.shape[0] + q = qkv[..., :q_lora_rank] + k_nope = qkv[..., q_lora_rank : q_lora_rank + kv_lora_rank] + q_contiguous = torch.empty_like(q) + if q.shape[0] > 0: + self.forward(input_q_a=q, input_kv_a=k_nope, output_q_a=q_contiguous) + + return qkv, q_contiguous, k_nope, None diff --git a/python/tokenspeed/runtime/layers/linear.py b/python/tokenspeed/runtime/layers/linear.py new file mode 100755 index 0000000..dc6f974 --- /dev/null +++ b/python/tokenspeed/runtime/layers/linear.py @@ -0,0 +1,1262 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + + +import torch +from torch.nn.parameter import Parameter + +from tokenspeed.runtime.distributed.comm_ops import all_gather, all_reduce +from tokenspeed.runtime.distributed.utils import divide, split_tensor_along_last_dim +from tokenspeed.runtime.layers.dense import ( + Fp8LinearMethod, + Mxfp4LinearMethod, + Nvfp4LinearMethod, + UnquantizedLinearMethod, + W8A8Fp8LinearMethod, +) +from tokenspeed.runtime.layers.parameter import ( + BaseWeightParameter, + BlockQuantScaleParameter, + PackedColumnParameter, + PackedWeightParameter, + PerTensorScaleParameter, + RowParallelWeightParameter, +) +from tokenspeed.runtime.layers.quantization import ( + Fp8Config, + Mxfp4Config, + Nvfp4Config, + W8A8Fp8Config, +) +from tokenspeed.runtime.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, +) +from tokenspeed.runtime.layers.quantization.compressed_tensors.compressed_tensors import ( + CompressedTensorsConfig, +) +from tokenspeed.runtime.layers.quantization.utils import ( + should_exclude_quant_module, + should_ignore_quant_layer, +) +from tokenspeed.runtime.utils import get_colorful_logger, set_weight_attrs + +logger = get_colorful_logger(__name__) + +WEIGHT_LOADER_V2_SUPPORTED = [ + "CompressedTensorsLinearMethod", + "AWQMarlinLinearMethod", + "AWQLinearMethod", + "GPTQMarlinLinearMethod", + "Fp8LinearMethod", + "BlockInt8LinearMethod", + "MarlinLinearMethod", + "QQQLinearMethod", + "GPTQMarlin24LinearMethod", + "TPUInt8LinearMethod", + "GPTQLinearMethod", + "IPEXAWQLinearMethod", +] + + +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, qkv_offsets: dict[str, tuple[int, int]], loaded_shard_id: str +) -> tuple[int, int]: + """Adjust the quantization offsets and sizes for BitsAndBytes sharding.""" + + total, _ = qkv_offsets["total"] + orig_offset, orig_size = qkv_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 + + +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. + override_kernel_name: Optional kernel name passed down to the + quant method's underlying ``tokenspeed_kernel.mm`` dispatch + (e.g. ``"cublaslt_mm_nvfp4"``). Lets the model force a + specific kernel for a particular layer. + interleave_linear_and_gate: If true, quantized post-load processing + prepares a 64-row linear/gate interleaved weight view for + fused GEMM+SwiGLU kernels. + """ + + def __init__( + self, + input_size: int, + output_size: int, + skip_bias_add: bool = False, + params_dtype: torch.dtype | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + override_kernel_name: str | None = None, + interleave_linear_and_gate: bool = False, + ): + 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 + self.prefix = prefix + self.override_kernel_name = override_kernel_name + self.interleave_linear_and_gate = interleave_linear_and_gate + + self.quant_config = quant_config + if quant_config is None or should_ignore_quant_layer( + prefix=prefix, ignored_layers=quant_config.ignored_layers + ): + self.quant_method: QuantizeMethodBase | None = UnquantizedLinearMethod() + elif isinstance(quant_config, Nvfp4Config): + # For NVFP4, excluded layers use unquantized (bf16) + if should_exclude_quant_module(prefix, quant_config.exclude_modules): + self.quant_method = UnquantizedLinearMethod() + else: + self.quant_method = Nvfp4LinearMethod(quant_config) + elif isinstance(quant_config, Mxfp4Config): + if getattr(quant_config, "use_dynamic_mxfp4_activations", False): + self.quant_method = Mxfp4LinearMethod(quant_config) + else: + # Existing MXFP4 support applies to MoE weights; dense weights + # remain unquantized unless the checkpoint stores dense MXFP4. + self.quant_method = UnquantizedLinearMethod() + elif isinstance(quant_config, CompressedTensorsConfig): + self.quant_method = quant_config.get_quant_method(self, prefix) + else: + if isinstance(quant_config, Fp8Config): + self.quant_method = Fp8LinearMethod(quant_config) + if isinstance(quant_config, W8A8Fp8Config): + self.quant_method = W8A8Fp8LinearMethod(quant_config) + + 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: torch.dtype | None = None, + quant_config: QuantizationConfig | None = 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( + layer=self, + input_size_per_partition=self.input_size, + 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, + ) + + 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, + shard_id=None, + begin_size=None, + ): + # 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) + + if begin_size is not None: + shard_size = loaded_weight.shape[0] + param[begin_size : begin_size + shard_size].data.copy_(loaded_weight) + elif shard_id is not None: + shard_size = loaded_weight.shape[0] + param[shard_id * shard_size : (shard_id + 1) * shard_size].data.copy_( + loaded_weight + ) + else: + assert param.size() == loaded_weight.size() + param.data.copy_(loaded_weight) + + def forward( + self, x: torch.Tensor, block_scale=None, output_dtype=None + ) -> torch.Tensor: + bias = self.bias if not self.skip_bias_add else None + assert self.quant_method is not None + if block_scale is not None: + # Note: block_scale is not None means flashinfer reduce-scatter fusion is used for fp8 block quant + # in this case, the input_ is already quantized to a fp8 tensor + output = self.quant_method.apply(self, x, bias, block_scale, output_dtype) + else: + 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: torch.dtype | None = None, + quant_config: QuantizationConfig | None = None, + output_sizes: list[int] | None = None, + prefix: str = "", + tp_rank: int | None = None, + tp_size: int | None = None, + tp_group: tuple[int, ...] | None = None, + use_presharded_weights: bool = False, + override_kernel_name: str | None = None, + interleave_linear_and_gate: bool = False, + ): + super().__init__( + input_size, + output_size, + skip_bias_add, + params_dtype, + quant_config, + prefix, + override_kernel_name, + interleave_linear_and_gate, + ) + self.gather_output = gather_output + self.use_presharded_weights = use_presharded_weights + assert self.quant_method is not None + + if tp_rank is None: + assert tp_size is None + assert tp_group is None + tp_rank, tp_size = 0, 1 + assert 0 <= tp_rank < tp_size + assert tp_size == 1 or tp_group is not None + self.tp_rank, self.tp_size, self.tp_group = tp_rank, tp_size, tp_group + + self.output_size_per_partition = divide(self.output_size, self.tp_size) + if output_sizes is None: + self.output_sizes = [self.output_size] + self.output_partition_sizes = [self.output_size_per_partition] + else: + self.output_sizes = output_sizes + # If QKV or MergedColumn, use output size of each partition. + self.output_partition_sizes = [ + divide(output_size, self.tp_size) for output_size in self.output_sizes + ] + + 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) + + 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 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) + + from tokenspeed.runtime.layers.parameter import _ColumnParallelWeightParameter + + if isinstance(param, _ColumnParallelWeightParameter): + param.load_column_parallel_weight( + loaded_weight, + tp_rank=self.tp_rank, + use_presharded_weights=self.use_presharded_weights, + ) + else: + # Some DeepSeek V3 AWQ checkpoints still reach the generic column + # loader path instead of the _ColumnParallelWeightParameter specialization. + param.load_column_parallel_weight(loaded_weight) + + def forward(self, input_, block_scale=None, output_dtype=None): + bias = self.bias if not self.skip_bias_add else None + + # Matrix multiply. + assert self.quant_method is not None + if block_scale is not None: + # Note: block_scale is not None means flashinfer all-reduce fusion is used for fp8 block quant + # in this case, the input_ is already quantized to a fp8 tensor + output_parallel = self.quant_method.apply( + self, input_, bias, block_scale, output_dtype + ) + else: + output_parallel = self.quant_method.apply(self, input_, bias) + if self.gather_output and self.tp_size > 1: + # All-gather across the partitions. + output = all_gather(output_parallel, self.tp_group, dim=-1) + 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: torch.dtype | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + tp_rank: int | None = None, + tp_size: int | None = None, + tp_group: tuple[int, ...] | None = None, + use_presharded_weights: bool = False, + override_kernel_name: str | None = None, + interleave_linear_and_gate: bool = False, + ): + 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, + output_sizes=output_sizes, + prefix=prefix, + tp_rank=tp_rank, + tp_size=tp_size, + tp_group=tp_group, + use_presharded_weights=use_presharded_weights, + override_kernel_name=override_kernel_name, + interleave_linear_and_gate=interleave_linear_and_gate, + ) + + def weight_loader( + self, + param: Parameter, + loaded_weight: torch.Tensor, + loaded_shard_id: int | None = None, + ): + 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) + 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 + ) + + 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 + # 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: BaseWeightParameter, 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, PackedWeightParameter)) + 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 + ) + # Special case for block-wise quantization scales. + # The scale tensor is smaller than the weight tensor by a factor + # of block_n, so we need to adjust offset and size accordingly. + elif isinstance(param, BlockQuantScaleParameter): + weight_block_size = self.quant_method.quant_config.weight_block_size + block_n = weight_block_size[0] + shard_offset = (shard_offset + block_n - 1) // block_n + shard_size = (shard_size + block_n - 1) // block_n + + 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: BaseWeightParameter, + loaded_weight: torch.Tensor, + loaded_shard_id: int | None = None, + ): + if loaded_shard_id is None: + if isinstance(param, PerTensorScaleParameter): + param.load_merged_column_weight(loaded_weight=loaded_weight, shard_id=0) + return + elif type(param) in (RowParallelWeightParameter, BaseWeightParameter): + param.load_merged_column_weight(loaded_weight=loaded_weight) + return + 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, + tp_rank=self.tp_rank, + shard_offset=shard_offset, + shard_size=shard_size, + use_presharded_weights=self.use_presharded_weights, + ) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + return None + + +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: int | None = None, + bias: bool = True, + skip_bias_add: bool = False, + params_dtype: torch.dtype | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + tp_rank: int | None = None, + tp_size: int | None = None, + tp_group: tuple[int, ...] | None = 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 + + tp_size = 1 if tp_size is None else 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 + input_size = self.hidden_size + output_size = ( + (self.num_heads + 2 * self.num_kv_heads) * tp_size * self.head_size + ) + 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 + ] + 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, + output_sizes=output_sizes, + prefix=prefix, + tp_rank=tp_rank, + tp_size=tp_size, + tp_group=tp_group, + use_presharded_weights=load_presharded_attn, + ) + + 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: BaseWeightParameter, 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, PackedWeightParameter)) + 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: BaseWeightParameter, + loaded_weight: torch.Tensor, + loaded_shard_id: str | None = 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 (RowParallelWeightParameter, BaseWeightParameter): + param.load_qkv_weight(loaded_weight=loaded_weight) + return + 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: str | None = None, + ): + 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) + 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 + + # 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) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + return None + + +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: torch.dtype | None = None, + reduce_results: bool = True, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + tp_rank: int | None = None, + tp_size: int | None = None, + tp_group: tuple[int, ...] | None = None, + use_presharded_weights: bool = False, + override_kernel_name: str | None = None, + interleave_linear_and_gate: bool = False, + ): + super().__init__( + input_size, + output_size, + skip_bias_add, + params_dtype, + quant_config, + prefix, + override_kernel_name, + interleave_linear_and_gate, + ) + + self.input_is_parallel = input_is_parallel + self.reduce_results = reduce_results + assert self.quant_method is not None + self.use_presharded_weights = use_presharded_weights + + if tp_rank is None: + assert tp_size is None + assert tp_group is None + tp_rank, tp_size = 0, 1 + assert 0 <= tp_rank < tp_size + assert tp_size == 1 or tp_group is not None + self.tp_rank, self.tp_size, self.tp_group = tp_rank, tp_size, tp_group + + self.input_size_per_partition = divide(input_size, self.tp_size) + + 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) + + 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 + 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 + ), f"{param_data.shape=} {loaded_weight.shape=}" + param_data.copy_(loaded_weight) + + def weight_loader_v2(self, param: BaseWeightParameter, 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, RowParallelWeightParameter): + # This `BaseWeightParameter` is defined in tokenspeed/runtime/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: + # Generic parameters do not support extra loading options. + param.load_row_parallel_weight(loaded_weight) + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + return None + + def forward(self, input_, scale=None): + 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 + + if scale is not None: + output_parallel = self.quant_method.apply( + self, input_parallel, bias_, scale, torch.bfloat16 + ) + else: + output_parallel = self.quant_method.apply(self, input_parallel, bias=bias_) + if self.reduce_results and self.tp_size > 1: + output = all_reduce(output_parallel, self.tp_group) + 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/tokenspeed/runtime/layers/logits_processor.py b/python/tokenspeed/runtime/layers/logits_processor.py new file mode 100755 index 0000000..7e63df6 --- /dev/null +++ b/python/tokenspeed/runtime/layers/logits_processor.py @@ -0,0 +1,766 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Logits processing.""" + +import dataclasses + +import torch +import triton +import triton.language as tl +from tokenspeed_kernel.ops.communication.triton import all_gather_inner, create_state +from tokenspeed_kernel.ops.sampling import argmax as sampling_argmax +from tokenspeed_kernel.ops.sampling.cute_dsl import ( + create_dist_argmax_state, + distributed_argmax, +) +from tokenspeed_kernel.platform import current_platform +from torch import nn + +from tokenspeed.runtime.distributed.comm_ops import all_gather_into_tensor +from tokenspeed.runtime.distributed.process_group_manager import ( + process_group_manager as pg_manager, +) +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.execution.forward_batch_info import ( + CaptureHiddenMode, + ForwardMode, +) +from tokenspeed.runtime.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding, +) +from tokenspeed.runtime.sampling.dp_sampling_config import ( + DpSamplingRuntimeConfig, +) +from tokenspeed.runtime.sampling.logits_layout import ( + LogitsLayoutExecutor, + LogitsLayoutPlan, +) +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +@dataclasses.dataclass +class LogitsProcessorOutput: + ## Part 1: This part will be assigned in python/tokenspeed/runtime/layers/logits_processor.py::LogitsProcessor + # The logits of the next tokens. shape: [#seq, vocab_size] + next_token_logits: torch.Tensor + # Used when ``do_argmax=True``. shape: [#seq] + next_token_ids: torch.Tensor | None = None + # Used by speculative decoding. + # The last hidden layers + hidden_states: torch.Tensor | None = None + logits_layout_plan: LogitsLayoutPlan | None = None + + ## Part 2: Populated by the active SamplingBackend during sample()/verify(). + # The logprobs of the next tokens. shape: [#seq] + next_token_logprobs: torch.Tensor | None = None + # The logprobs and ids of the top-k tokens in output positions. shape: [#seq, k] + next_token_top_logprobs_val: list | None = None + next_token_top_logprobs_idx: list | None = 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: list | None = None + next_token_token_ids_logprobs_idx: list | None = None + + ## Part 3: Prefill-only. This part will be assigned in python/tokenspeed/runtime/layers/logits_processor.py::LogitsProcessor + # The logprobs of input tokens. shape: [#token] + input_token_logprobs: torch.Tensor | None = 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: list | None = None + input_token_ids_logprobs_idx: list | None = None + + +@dataclasses.dataclass +class LogitsMetadata: + forward_mode: ForwardMode + capture_hidden_mode: CaptureHiddenMode = CaptureHiddenMode.NULL + gather_ids: torch.Tensor | None = None + + extend_return_logprob: bool = False + extend_return_top_logprob: bool = False + extend_token_ids_logprob: bool = False + extend_seq_lens_cpu: list[int] | None = None + extend_logprob_start_lens_cpu: list[int] | None = None + extend_logprob_pruned_lens_cpu: list[int] | None = None + top_logprobs_nums: list[int] | None = None + extend_input_logprob_token_ids_gpu: torch.Tensor | None = None + token_ids_logprobs: list[list[int]] | None = 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: torch.Tensor | None = None + # The start position of local hidden states. + dp_local_start_pos: torch.Tensor | None = None + dp_local_num_tokens: torch.Tensor | None = None + gathered_buffer: torch.Tensor | None = None + # Buffer to gather logits from all ranks. + forward_batch_gathered_buffer: torch.Tensor | None = None + + @classmethod + def from_forward_context(cls, ctx: ForwardContext): + return cls( + forward_mode=ctx.forward_mode, + capture_hidden_mode=ctx.capture_hidden_mode, + gather_ids=ctx.gather_ids, + ) + + +_FUSED_LM_HEAD_GEMM = None + + +def _get_fused_lm_head_gemm(): + """Lazily import the fused lm_head GEMM kernel. + + The kernel is only present when tokenspeed-kernel was built with a + compatible nvcc. Cache a sentinel when unavailable so we fall back + to ``torch.matmul`` silently on subsequent calls. + """ + global _FUSED_LM_HEAD_GEMM + if _FUSED_LM_HEAD_GEMM is not None: + return _FUSED_LM_HEAD_GEMM + if not current_platform().is_nvidia: + _FUSED_LM_HEAD_GEMM = (None, None) + return _FUSED_LM_HEAD_GEMM + try: + from tokenspeed_kernel.thirdparty.cuda.lm_head_gemm import ( + lm_head_gemm, + should_use_fused, + ) + + _FUSED_LM_HEAD_GEMM = (should_use_fused, lm_head_gemm) + except Exception: + _FUSED_LM_HEAD_GEMM = (None, None) + return _FUSED_LM_HEAD_GEMM + + +def _lm_head_matmul(hidden_states: torch.Tensor, weight: torch.Tensor) -> torch.Tensor: + """Compute ``hidden_states @ weight.T``. + + Routes to the fused ``lm_head_gemm`` when the shape matches a compiled + template and the bench-driven perf gate accepts (``should_use_fused``). + Otherwise falls back to ``torch.matmul``. + + Only enabled for Kimi (``model_type == "kimi_k2"``) at the call site — + on DSv3 the fused kernel's PDL launch surface caused a downstream EAGLE3 + spec decode AR regression that we have not characterised end-to-end; on + Kimi the perf win is the largest and the regression has not been + reproduced, so we gate the fused path to Kimi only. + """ + cast_hidden = hidden_states.to(weight.dtype) + should_use_fused, lm_head_gemm = _get_fused_lm_head_gemm() + if should_use_fused is not None and should_use_fused(cast_hidden, weight): + return lm_head_gemm(cast_hidden, weight, enable_pdl=True) + return torch.matmul(cast_hidden, weight.T) + + +class LogitsProcessor(nn.Module): + + _LOGITS_AG_MAX_TOKENS = 128 + _LOGITS_AG_STATE_UNINITIALIZED = object() + _LOGITS_AG_STATES = {} + + _LOGITS_DIST_ARGMAX_MAX_TOKENS = 8192 + _LOGITS_DIST_ARGMAX_UNINITIALIZED = object() + _LOGITS_DIST_ARGMAX_STATES = {} + + def __init__( + self, + config, + skip_all_gather: bool = False, + do_argmax: bool = False, + logit_scale: float | None = None, + tp_rank: int | None = None, + tp_size: int | None = None, + tp_group: tuple[int, ...] | None = None, + ): + super().__init__() + self.config = config + self.skip_all_gather = skip_all_gather + self.do_argmax = do_argmax + self.dp_sampling_enabled = False + self.dp_num_tokens_per_req = 1 + self.dp_sampling_min_bs = 0 + self.logit_scale = logit_scale + self._logits_layout_executor: LogitsLayoutExecutor | None = None + + if tp_rank is None: + if tp_size is not None or tp_group is not None: + raise ValueError("tp_size and tp_group require tp_rank.") + tp_rank, tp_size = 0, 1 + elif tp_size is None: + raise ValueError("tp_size is required when tp_rank is provided.") + if not 0 <= tp_rank < tp_size: + raise ValueError(f"Invalid tensor-parallel rank: {tp_rank}/{tp_size}.") + if tp_size != 1 and tp_group is None: + raise ValueError("tp_group is required when tp_size > 1.") + self.tp_rank, self.tp_size, self.tp_group = tp_rank, tp_size, tp_group + + self._all_gather_state = self._LOGITS_AG_STATE_UNINITIALIZED + self._dist_argmax_state = self._LOGITS_DIST_ARGMAX_UNINITIALIZED + + 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 + + # Gate the fused lm_head GEMM to Kimi only. See ``_lm_head_matmul``. + self._use_fused_lm_head = getattr(self.config, "model_type", None) == "kimi_k2" + + def configure_dp_logits_layout(self, runtime: DpSamplingRuntimeConfig) -> None: + if ( + not runtime.enabled + or runtime.topology is None + or runtime.min_bs is None + or runtime.max_bucket_bs is None + or runtime.vocab_size is None + or runtime.device is None + ): + raise RuntimeError("enabled DP sampling runtime is incomplete") + topology = runtime.topology + self.dp_sampling_enabled = True + self.dp_num_tokens_per_req = runtime.num_tokens_per_req + self.dp_sampling_min_bs = runtime.min_bs + self._logits_layout_executor = LogitsLayoutExecutor( + tp_rank=topology.tp_rank, + tp_size=topology.tp_size, + tp_group=topology.tp_group, + max_bucket_bs=runtime.max_bucket_bs, + num_tokens_per_req=runtime.num_tokens_per_req, + vocab_size=runtime.vocab_size, + device=runtime.device, + ) + + def _resolve_logits_layout_plan( + self, + hidden_states: torch.Tensor, + logits_metadata: LogitsMetadata, + ) -> LogitsLayoutPlan | None: + if not self.dp_sampling_enabled: + return None + if not logits_metadata.forward_mode.is_decode(): + return None + n = self.dp_num_tokens_per_req + rows = hidden_states.shape[0] + if rows % n != 0: + raise ValueError(f"hidden_states have {rows} rows, not divisible by N={n}") + effective_bs = rows // n + bucket_bs = ((effective_bs + self.tp_size - 1) // self.tp_size) * self.tp_size + if effective_bs < self.dp_sampling_min_bs: + return None + return LogitsLayoutPlan( + effective_bs=effective_bs, + bucket_bs=bucket_bs, + tp_size=self.tp_size, + num_tokens_per_req=n, + ) + + def _init_all_gather_state(self, lm_head: VocabParallelEmbedding): + if not current_platform().is_nvidia: + return None + + if self.tp_size == 1 or self.skip_all_gather: + return None + + vocab_padded = lm_head.weight.size(0) * self.tp_size + if vocab_padded % (self.tp_size * 8) != 0: + return None + + key = (self.tp_group, vocab_padded) + if key not in self._LOGITS_AG_STATES: + self._LOGITS_AG_STATES[key] = create_state( + group=pg_manager.get_process_group("nccl", self.tp_group), + rank_in_group=self.tp_rank, + max_tokens=self._LOGITS_AG_MAX_TOKENS, + hidden_size=vocab_padded, + ) + return self._LOGITS_AG_STATES[key] + + def _init_dist_argmax_state(self, lm_head: VocabParallelEmbedding): + if not current_platform().is_nvidia: + return None + + if self.tp_size == 1 or self.skip_all_gather or self.dp_sampling_enabled: + return None + + vocab_per_rank = lm_head.weight.size(0) + if vocab_per_rank * self.tp_size != self.config.vocab_size: + return None # padded vocab: sharded argmax could pick a pad column + if vocab_per_rank < 4096 or vocab_per_rank % 32 != 0: + return None # below the kernel's vocab floor / alignment + + key = (self.tp_group, vocab_per_rank) + if key not in self._LOGITS_DIST_ARGMAX_STATES: + self._LOGITS_DIST_ARGMAX_STATES[key] = create_dist_argmax_state( + group=pg_manager.get_process_group("nccl", self.tp_group), + rank_in_group=self.tp_rank, + max_M=self._LOGITS_DIST_ARGMAX_MAX_TOKENS, + dtype=lm_head.weight.dtype, + device=lm_head.weight.device, + skip_ping_pong=True, + ) + + return self._LOGITS_DIST_ARGMAX_STATES[key] + + def forward( + self, + input_ids, + hidden_states, + lm_head: VocabParallelEmbedding, + logits_metadata: LogitsMetadata, + aux_hidden_states: torch.Tensor | None = None, + ) -> LogitsProcessorOutput: + # Get the last hidden states and last logits for the next token prediction + if not logits_metadata.extend_return_logprob: + gather_ids = logits_metadata.gather_ids + if gather_ids is not None: + # Shapes align iff midlayer already pruned to one row per request + # (draft first-step reduce). Other paths emit [N, H] with N > bs. + if gather_ids.shape[0] == hidden_states.shape[0]: + pruned_states = hidden_states + if aux_hidden_states is not None: + aux_pruned_states = list(aux_hidden_states) + else: + pruned_states = hidden_states[gather_ids] + if aux_hidden_states is not None: + aux_pruned_states = [h[gather_ids] for h in aux_hidden_states] + else: + if logits_metadata.forward_mode.is_extend_or_mixed(): + raise RuntimeError( + "EXTEND/MIXED forward must set gather_ids on ForwardContext" + ) + pruned_states = hidden_states + if aux_hidden_states is not None: + aux_pruned_states = list(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. + if extend_len <= start_len: + raise RuntimeError("extend_len must be greater than 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_layout_plan = self._resolve_logits_layout_plan( + pruned_states, logits_metadata + ) + logits = self._get_logits( + pruned_states, lm_head, logits_metadata, plan=logits_layout_plan + ) + sampled_logits = ( + logits[sample_indices] if sample_indices is not None else logits + ) + + hidden_states_to_store: torch.Tensor | None = 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 = ( + aux_hidden_states[0] + if len(aux_hidden_states) == 1 + else 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 = ( + aux_pruned_states[0] + if len(aux_pruned_states) == 1 + else 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: + raise RuntimeError("Should never reach") + + if not logits_metadata.extend_return_logprob: + # Decode mode or extend mode without return_logprob. + # Greedy draft path: emit token ids here, fusing the cross-rank + # vocab reduction into the argmax when gated on. + next_token_ids = self._argmax(sampled_logits) if self.do_argmax else None + return LogitsProcessorOutput( + next_token_logits=sampled_logits, + next_token_ids=next_token_ids, + hidden_states=hidden_states_to_store, + logits_layout_plan=logits_layout_plan, + ) + 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: torch.Tensor | None = None, + plan: LogitsLayoutPlan | None = 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. + """ + dp_sampling = plan is not None + if dp_sampling and not self.dp_sampling_enabled: + raise RuntimeError( + "DP logits layout plan was provided but LogitsProcessor was not " + "configured with dp_sampling" + ) + + if dp_sampling and self.skip_all_gather: + if self._logits_layout_executor is None: + raise RuntimeError( + "dp_sampling logits layout executor is not configured" + ) + hidden_states = self._logits_layout_executor.slice_hidden_states( + hidden_states, plan + ) + + if hasattr(lm_head, "weight"): + if self._use_fused_lm_head: + logits = _lm_head_matmul(hidden_states, lm_head.weight) + else: + logits = torch.matmul( + hidden_states.to(lm_head.weight.dtype), lm_head.weight.T + ) + else: + # GGUF models + logits = lm_head.linear_method.apply(lm_head, hidden_states, embedding_bias) + + if self.logit_scale is not None: + logits.mul_(self.logit_scale) + + if dp_sampling and not self.skip_all_gather: + if self._logits_layout_executor is None: + raise RuntimeError( + "dp_sampling logits layout executor is not configured" + ) + logits = self._logits_layout_executor.swap_batch_vocab(logits, plan) + + elif not dp_sampling and self.tp_size > 1 and not self.skip_all_gather: + if self.do_argmax: + if self._dist_argmax_state is self._LOGITS_DIST_ARGMAX_UNINITIALIZED: + self._dist_argmax_state = self._init_dist_argmax_state(lm_head) + + if ( + self._dist_argmax_state is not None + and not self.final_logit_softcapping + and logits.size(0) <= self._LOGITS_DIST_ARGMAX_MAX_TOKENS + ): + return logits + + if self._all_gather_state is self._LOGITS_AG_STATE_UNINITIALIZED: + self._all_gather_state = self._init_all_gather_state(lm_head) + + if ( + self._all_gather_state is not None + and logits.size(0) <= self._LOGITS_AG_MAX_TOKENS + ): + # skip_entry_sync=True assumes other sync points existing between two all_gather_inner calls. + logits = all_gather_inner( + self._all_gather_state, + logits, + tp_hidden_dim=logits.size(-1) * self.tp_size, + skip_entry_sync=True, + safe=False, + ) + else: + num_rows = logits.size(0) + local_vocab_size = logits.size(1) + gathered_logits = torch.empty( + self.tp_size * num_rows, + local_vocab_size, + dtype=logits.dtype, + device=logits.device, + ) + all_gather_into_tensor(gathered_logits, logits, self.tp_group) + logits = ( + gathered_logits.view(self.tp_size, num_rows, local_vocab_size) + .transpose(0, 1) + .contiguous() + .view(num_rows, local_vocab_size * self.tp_size) + ) + + logits = logits[:, : self.config.vocab_size].contiguous() + + if self.final_logit_softcapping: + fused_softcap_generic(logits, self.final_logit_softcapping) + + return logits + + def _argmax(self, logits: torch.Tensor) -> torch.Tensor: + if ( + self._dist_argmax_state + not in (self._LOGITS_DIST_ARGMAX_UNINITIALIZED, None) + and not self.final_logit_softcapping + and logits.size(0) <= self._LOGITS_DIST_ARGMAX_MAX_TOKENS + ): + _, idx = distributed_argmax(self._dist_argmax_state, logits) + return idx + else: + return sampling_argmax(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 + """ + last_logits = last_logits.float() + # 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 + # 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 tokenspeed.runtime.sampling.utils 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).to(tl.float32) + + # Perform operations in-place + x = x / softcapping_value + + # Stable tanh form; the exp ratio overflows to inf/inf for large logits. + x = 2 * tl.sigmoid(2 * x) - 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 + + +def fused_softcap_generic(full_logits, final_logit_softcapping): + return fused_softcap(full_logits, final_logit_softcapping) diff --git a/python/tokenspeed/runtime/layers/moe/__init__.py b/python/tokenspeed/runtime/layers/moe/__init__.py new file mode 100755 index 0000000..68a877b --- /dev/null +++ b/python/tokenspeed/runtime/layers/moe/__init__.py @@ -0,0 +1,35 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from tokenspeed.runtime.layers.moe.expert import MoELayer +from tokenspeed.runtime.layers.moe.loader import ( + MoECheckpointLoader, + MoECheckpointLoadError, + build_moe_checkpoint_loader, +) +from tokenspeed.runtime.layers.moe.schema import ExpertCheckpointSchema + +__all__ = [ + "ExpertCheckpointSchema", + "MoECheckpointLoadError", + "MoECheckpointLoader", + "MoELayer", + "build_moe_checkpoint_loader", +] diff --git a/python/tokenspeed/runtime/layers/moe/config.py b/python/tokenspeed/runtime/layers/moe/config.py new file mode 100755 index 0000000..25460ec --- /dev/null +++ b/python/tokenspeed/runtime/layers/moe/config.py @@ -0,0 +1,39 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from dataclasses import dataclass + +import torch + + +@dataclass +class ExpertParallelConfig: + top_k: int + num_experts: int + low_latency_max_num_tokens_per_gpu: int + max_num_tokens_per_gpu: int + hidden_size: int + rank: int + world_size: int + group: torch.distributed.ProcessGroup + params_dtype: torch.dtype + + +EPConfig = ExpertParallelConfig diff --git a/python/tokenspeed/runtime/layers/moe/expert.py b/python/tokenspeed/runtime/layers/moe/expert.py new file mode 100755 index 0000000..2ebe22c --- /dev/null +++ b/python/tokenspeed/runtime/layers/moe/expert.py @@ -0,0 +1,277 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + + +import tokenspeed_kernel +import torch + +from tokenspeed.runtime.distributed.process_group_manager import ( + process_group_manager as pg_manager, +) +from tokenspeed.runtime.layers.activation import SwigluArg +from tokenspeed.runtime.layers.moe.topk import TopKOutput, TopKOutputFormat +from tokenspeed.runtime.layers.moe.types import MoELayerSpec +from tokenspeed.runtime.layers.moe.utils import ( + RoutingMethodType, + get_all2all_backend, + get_moe_backend, +) +from tokenspeed.runtime.layers.moe.weights import create_layer_weights +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.layers.quantization.utils import ( + should_exclude_quant_module, + should_ignore_quant_layer, +) +from tokenspeed.runtime.utils.env import global_server_args_dict +from tokenspeed.runtime.utils.pdl import pdl_enabled + + +class MoELayer(torch.nn.Module): + def __init__( + self, + top_k: int, + num_experts: int, + hidden_size: int, + intermediate_size: int, + quant_config: QuantizationConfig, + layer_index: int, + prefix: str = "", + tp_rank: int | None = None, + tp_size: int | None = None, + ep_rank: int | None = None, + ep_size: int | None = None, + zero_expert_type: str = "", + activation: str = "silu", + activation_alpha=None, + swiglu_limit=None, + swiglu_beta: float | None = None, + w13_input_layout: str = "concatenated", + with_bias=False, + routing_config: dict = {}, + ): + super().__init__() + self.layer_index = layer_index + self.prefix = prefix + self.top_k = top_k + self.num_experts = num_experts + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.quant_config = quant_config + self.ep_num_redundant_experts = global_server_args_dict[ + "ep_num_redundant_experts" + ] + self.zero_expert_type = zero_expert_type + self.activation = activation + self.swiglu_arg = None + if self.activation == "swiglu": + self.swiglu_arg = SwigluArg(alpha=activation_alpha, limit=swiglu_limit) + # Per-model knobs the MoE backend reads in process_weights_after_loading. + # ``swiglu_beta``: gpt-oss uses silu(α·gate)·(up + 1) and sets 1.0; + # standard SwiGLU (e.g. deepseek-v4) leaves it None. + # ``w13_input_layout``: "interleaved" for HF gpt-oss-style row layout + # ([w1_0, w3_0, w1_1, w3_1, ...]); "concatenated" (default) for the + # shared MoE checkpoint loader's [w1_all | w3_all] block layout. + self.swiglu_beta = swiglu_beta + if w13_input_layout not in {"interleaved", "concatenated"}: + raise ValueError( + f"w13_input_layout must be 'interleaved' or 'concatenated', " + f"got {w13_input_layout!r}" + ) + self.w13_input_layout = w13_input_layout + + if tp_rank is None: + assert tp_size is None + tp_rank, tp_size = 0, 1 + self.tp_rank, self.tp_size = tp_rank, tp_size + self.moe_tp_size = self.tp_size + if ep_rank is None: + assert ep_size is None + ep_rank, ep_size = 0, 1 + self.ep_rank, self.ep_size = ep_rank, ep_size + + if tp_size > 1 and ep_size > 1: + raise ValueError("Mixed TP and EP is not supported yet.") + + num_local_experts = num_experts // self.ep_size + + self.num_local_experts = num_local_experts + self._spec = MoELayerSpec( + top_k=top_k, + num_experts=num_experts, + num_local_experts=num_local_experts, + hidden_size=hidden_size, + intermediate_size=intermediate_size, + activation=activation, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + ep_rank=self.ep_rank, + ep_size=self.ep_size, + prefix=prefix, + a2a_backend=get_all2all_backend().value, + ) + + # Routing config + self.routing_config = routing_config + self._correction_bias = routing_config.get("correction_bias", None) + self._routing_method_type = routing_config.get( + "routing_method_type", RoutingMethodType.DeepSeekV3 + ) + self._routing_logits_dtype = torch.bfloat16 + if self._routing_method_type in ( + RoutingMethodType.DeepSeekV3, + RoutingMethodType.MiniMax2, + ): + self._routing_logits_dtype = torch.float32 + self._n_group = routing_config.get("n_group", 0) + self._topk_group = routing_config.get("topk_group", 0) + self._routed_scaling_factor = routing_config.get("routed_scaling_factor", 1.0) + self._normalize_topk_weights = routing_config.get( + "normalize_topk_weights", True + ) + + # Quantization config. ignored_layers (compressed-tensors) keys the MoE + # block; exclude_modules (ModelOpt) keys the fused experts. + self._quant_kind = "unquant" + if ( + quant_config is not None + and not should_ignore_quant_layer(self.prefix, quant_config.ignored_layers) + and not should_exclude_quant_module( + f"{self.prefix}.experts", quant_config.exclude_modules + ) + ): + self._quant_kind = quant_config.moe_weight_dtype() + + fp8_scale_block_shape = None + internal_activation_dtype = "input" + if self._quant_kind == "fp8": + fp8_scale_block_shape = tuple(self.quant_config.weight_block_size) + if self._quant_kind == "mxfp4": + if self.quant_config.is_w4a8_fp8: + internal_activation_dtype = "fp8" + elif getattr(self.quant_config, "use_dynamic_mxfp4_activations", False): + internal_activation_dtype = "input" + + input_dtype = torch.get_default_dtype() + if input_dtype not in {torch.float16, torch.bfloat16}: + input_dtype = torch.float16 + + deepep_group = None + if self._spec.use_deepep: + mapping = global_server_args_dict["mapping"] + deepep_group = pg_manager.get_process_group( + "nccl", + mapping.moe.tp_ep_group, + ) + + # Moe Backend plan + moe_backend = get_moe_backend().value + moe_backend = None if moe_backend == "auto" else moe_backend + self.plan = tokenspeed_kernel.moe_plan( + self._quant_kind, + input_dtype=input_dtype, + activation=self.activation, + a2a_backend=self._spec.a2a_backend, + ep_size=self.ep_size, + ispp=self.intermediate_size // self.tp_size, + fp8_scale_block_shape=fp8_scale_block_shape, + internal_activation_dtype=internal_activation_dtype, + with_bias=with_bias, + deepep_group=deepep_group, + solution=moe_backend, + ) + + create_layer_weights( + self._spec, + self, + self._quant_kind, + self.quant_config, + with_bias=with_bias, + solution=self.plan["solution"], + ) + self._weights_processed = False + + def process_weights_after_loading(self, module) -> None: + if self._weights_processed: + return + + tokenspeed_kernel.moe_process_weights(self.plan, module) + self._weights_processed = True + + @property + def support_routing(self) -> bool: + return self.plan["support_routing"] + + @property + def topk_output_format(self): + if self.support_routing: + return TopKOutputFormat.BYPASSED + return TopKOutputFormat.STANDARD + + @property + def supports_deferred_finalize(self) -> bool: + return self.plan["supports_deferred_finalize"] + + def forward_zero_experts(self, topk_output): + zero_expert_limit = self.num_experts + if self.ep_num_redundant_experts is not None: + zero_expert_limit = zero_expert_limit - self.ep_num_redundant_experts + + normal_expert_mask = topk_output.topk_ids >= zero_expert_limit + topk_output.topk_ids[normal_expert_mask] = -1 + if self.zero_expert_type == "copy": + topk_output.topk_weights[normal_expert_mask] = 1.0 + if self.zero_expert_type == "drop": + topk_output.topk_weights[normal_expert_mask] = 0.0 + + def forward( + self, + hidden_states: torch.Tensor, + topk_output: TopKOutput, + num_global_tokens: int, + max_num_tokens_per_gpu: int, + do_finalize: bool = True, + ): + if not do_finalize and not self.supports_deferred_finalize: + raise AssertionError("MoELayer does not support do_finalize=False") + + if self.support_routing: + return tokenspeed_kernel.moe_apply( + self.plan, + hidden_states, + self, + topk_output.router_logits, + num_tokens_global=num_global_tokens, + max_num_tokens_per_gpu=max_num_tokens_per_gpu, + do_finalize=do_finalize, + enable_pdl=pdl_enabled(), + ) + else: + return tokenspeed_kernel.moe_apply( + self.plan, + hidden_states, + self, + topk_output.router_logits, + topk_weights=topk_output.topk_weights, + topk_ids=topk_output.topk_ids, + num_tokens_global=num_global_tokens, + max_num_tokens_per_gpu=max_num_tokens_per_gpu, + do_finalize=do_finalize, + enable_pdl=pdl_enabled(), + ) diff --git a/python/tokenspeed/runtime/layers/moe/loader.py b/python/tokenspeed/runtime/layers/moe/loader.py new file mode 100644 index 0000000..0213a8a --- /dev/null +++ b/python/tokenspeed/runtime/layers/moe/loader.py @@ -0,0 +1,479 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass + +import torch +from torch import nn + +from tokenspeed.runtime.layers.moe.schema import ExpertCheckpointSchema +from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader + + +@dataclass(frozen=True) +class CheckpointPlanEntry: + param_name: str + checkpoint_weight_name: str + shard_id: str + + def matches(self, checkpoint_name: str) -> bool: + return self.checkpoint_weight_name in checkpoint_name + + def resolve_param_name(self, checkpoint_name: str) -> str: + return checkpoint_name.replace(self.checkpoint_weight_name, self.param_name) + + +@dataclass(frozen=True) +class ExpertWeightPlanEntry(CheckpointPlanEntry): + local_expert_id: int + + +@dataclass(frozen=True) +class FusedExpertWeightPlanEntry(CheckpointPlanEntry): + split_dim: int | None = None + split_chunks: int | None = None + split_index: int | None = None + + +class MoECheckpointLoadError(RuntimeError): + pass + + +def _build_default_expert_plan( + schema: ExpertCheckpointSchema, + *, + num_experts: int, + ep_rank: int, + ep_size: int, +) -> list[ExpertWeightPlanEntry]: + # Expert ownership is assumed to be a contiguous per-rank range here. + # EPLB-aware remapping would need a different planning step. + num_local_experts = num_experts // ep_size + start_expert = num_local_experts * ep_rank + expert_plan: list[ExpertWeightPlanEntry] = [] + for local_expert_id in range(num_local_experts): + expert_id = start_expert + local_expert_id + expert_plan.extend( + ( + ExpertWeightPlanEntry( + param_name="experts.w13_", + checkpoint_weight_name=schema.make_expert_weight_name( + expert_id, "gate_proj" + ), + shard_id="w1", + local_expert_id=local_expert_id, + ), + ExpertWeightPlanEntry( + param_name="experts.w13_", + checkpoint_weight_name=schema.make_expert_weight_name( + expert_id, "up_proj" + ), + shard_id="w3", + local_expert_id=local_expert_id, + ), + ExpertWeightPlanEntry( + param_name="experts.w2_", + checkpoint_weight_name=schema.make_expert_weight_name( + expert_id, "down_proj" + ), + shard_id="w2", + local_expert_id=local_expert_id, + ), + ) + ) + return expert_plan + + +def _build_global_expert_name_plan( + schema: ExpertCheckpointSchema, + *, + num_experts: int, +) -> list[CheckpointPlanEntry]: + expert_plan: list[CheckpointPlanEntry] = [] + for expert_id in range(num_experts): + expert_plan.extend( + ( + CheckpointPlanEntry( + param_name="", + checkpoint_weight_name=schema.make_expert_weight_name( + expert_id, "gate_proj" + ), + shard_id="", + ), + CheckpointPlanEntry( + param_name="", + checkpoint_weight_name=schema.make_expert_weight_name( + expert_id, "up_proj" + ), + shard_id="", + ), + CheckpointPlanEntry( + param_name="", + checkpoint_weight_name=schema.make_expert_weight_name( + expert_id, "down_proj" + ), + shard_id="", + ), + ) + ) + return expert_plan + + +def _build_default_fused_plan( + schema: ExpertCheckpointSchema, + *, + fused_gate_up_as_w13: bool = False, + include_bias: bool = False, +) -> list[FusedExpertWeightPlanEntry]: + if fused_gate_up_as_w13: + fused_plan = [ + FusedExpertWeightPlanEntry( + param_name="experts.w13_weight", + checkpoint_weight_name=( + f"experts.{schema.get_semantic_name('gate_up_fused')}" + ), + shard_id="w13", + ), + FusedExpertWeightPlanEntry( + param_name="experts.w2_weight", + checkpoint_weight_name=f"experts.{schema.get_semantic_name('down_proj')}", + shard_id="w2", + ), + ] + if include_bias: + fused_plan.extend( + ( + FusedExpertWeightPlanEntry( + param_name="experts.w13_weight_bias", + checkpoint_weight_name=( + f"experts.{schema.get_semantic_name('gate_up_bias')}" + ), + shard_id="w13", + ), + FusedExpertWeightPlanEntry( + param_name="experts.w2_weight_bias", + checkpoint_weight_name=( + f"experts.{schema.get_semantic_name('down_bias')}" + ), + shard_id="w2", + ), + ) + ) + return fused_plan + + fused_plan = [ + FusedExpertWeightPlanEntry( + param_name="experts.w13_weight", + checkpoint_weight_name=f"experts.{schema.get_semantic_name('gate_up_fused')}", + shard_id="w1", + split_dim=-2, + split_chunks=2, + split_index=0, + ), + FusedExpertWeightPlanEntry( + param_name="experts.w13_weight", + checkpoint_weight_name=f"experts.{schema.get_semantic_name('gate_up_fused')}", + shard_id="w3", + split_dim=-2, + split_chunks=2, + split_index=1, + ), + FusedExpertWeightPlanEntry( + param_name="experts.w2_weight", + checkpoint_weight_name=f"experts.{schema.get_semantic_name('down_proj')}", + shard_id="w2", + ), + ] + if include_bias: + fused_plan.extend( + ( + FusedExpertWeightPlanEntry( + param_name="experts.w13_weight_bias", + checkpoint_weight_name=( + f"experts.{schema.get_semantic_name('gate_up_bias')}" + ), + shard_id="w1", + split_dim=-1, + split_chunks=2, + split_index=0, + ), + FusedExpertWeightPlanEntry( + param_name="experts.w13_weight_bias", + checkpoint_weight_name=( + f"experts.{schema.get_semantic_name('gate_up_bias')}" + ), + shard_id="w3", + split_dim=-1, + split_chunks=2, + split_index=1, + ), + FusedExpertWeightPlanEntry( + param_name="experts.w2_weight_bias", + checkpoint_weight_name=f"experts.{schema.get_semantic_name('down_bias')}", + shard_id="w2", + ), + ) + ) + return fused_plan + + +def _load_fused_expert_tensor( + param, + loaded_weight, + *, + shard_id: str, + num_experts: int, + ep_rank: int, + ep_size: int, +) -> None: + # Expert ownership is assumed to be a contiguous per-rank range here. + # EPLB-aware remapping would need a different loading step. + num_local_experts = num_experts // ep_size + start_expert = num_local_experts * ep_rank + end_expert = start_expert + num_local_experts + weight_loader = param.weight_loader + for expert_id in range(start_expert, end_expert): + local_expert_id = expert_id - start_expert + weight_loader( + param, + loaded_weight[expert_id], + shard_id=shard_id, + local_expert_id=local_expert_id, + ) + + +class MoECheckpointLoader: + def __init__( + self, + *, + params_dict: dict[str, nn.Parameter], + expert_plan: Sequence[ExpertWeightPlanEntry] = (), + global_expert_plan: Sequence[CheckpointPlanEntry] = (), + fused_plan: Sequence[FusedExpertWeightPlanEntry] = (), + num_experts: int | None = None, + ep_rank: int = 0, + ep_size: int = 1, + fused_load_style: str = "per_expert", + transpose_local_tensor_non_bias: bool = False, + ) -> None: + self._params_dict = params_dict + self._expert_plan = tuple(expert_plan) + self._global_expert_plan = tuple(global_expert_plan) + self._fused_plan = tuple(fused_plan) + self._num_experts = num_experts + self._ep_rank = ep_rank + self._ep_size = ep_size + self._fused_load_style = fused_load_style + self._transpose_local_tensor_non_bias = transpose_local_tensor_non_bias + + if self._fused_plan and self._num_experts is None: + raise ValueError("num_experts is required when fused_plan is used") + if fused_load_style not in {"per_expert", "local_tensor"}: + raise ValueError(f"Unknown fused_load_style: {fused_load_style}") + + @staticmethod + def _matches_plan(plan: Sequence[CheckpointPlanEntry], name: str) -> bool: + return any(plan_entry.matches(name) for plan_entry in plan) + + def matches(self, name: str) -> bool: + return self._matches_plan(self._fused_plan, name) or self._matches_plan( + self._expert_plan, name + ) + + def is_expert_checkpoint_weight(self, name: str) -> bool: + """Return whether ``name`` belongs to this loader's MoE checkpoint schema. + + Args: + name: Checkpoint tensor name after any model-specific remapping. + + Returns: + ``True`` for local, non-local, or fused expert checkpoint tensors + that this loader is responsible for; ``False`` for unrelated + checkpoint tensors. + """ + return self._matches_plan(self._fused_plan, name) or self._matches_plan( + self._global_expert_plan, name + ) + + def _load_expert(self, name: str, loaded_weight: torch.Tensor) -> str | None: + mapped_name: str | None = None + for plan_entry in self._expert_plan: + if not plan_entry.matches(name): + continue + + mapped_name = plan_entry.resolve_param_name(name) + param = self._params_dict.get(mapped_name) + if param is None: + continue + + param.weight_loader( + param, + loaded_weight, + shard_id=plan_entry.shard_id, + local_expert_id=plan_entry.local_expert_id, + ) + return mapped_name + + if mapped_name is not None: + self._raise_unloaded_match(name, mapped_name) + return None + + @staticmethod + def _raise_unloaded_match(name: str, mapped_name: str | None) -> None: + if mapped_name is None: + raise MoECheckpointLoadError( + f"Matched MoE checkpoint mapping for {name!r} but did not load any parameter" + ) + raise MoECheckpointLoadError( + f"Matched MoE checkpoint mapping for {name!r} -> {mapped_name!r}, " + "but the target parameter was not found or no tensor was loaded" + ) + + @staticmethod + def _raise_unmatched(name: str) -> None: + raise MoECheckpointLoadError( + f"{name!r} does not match any MoE checkpoint mapping" + ) + + def _load_fused(self, name: str, loaded_weight: torch.Tensor) -> str | None: + matched_entries = [ + plan_entry for plan_entry in self._fused_plan if plan_entry.matches(name) + ] + if not matched_entries: + return None + + selected_checkpoint_weight_name = max( + (plan_entry.checkpoint_weight_name for plan_entry in matched_entries), + key=len, + ) + + loaded_any = False + mapped_name: str | None = None + + for plan_entry in matched_entries: + if plan_entry.checkpoint_weight_name != selected_checkpoint_weight_name: + continue + + mapped_name = plan_entry.resolve_param_name(name) + param = self._params_dict.get(mapped_name) + if param is None: + continue + + tensor_to_load = loaded_weight + if plan_entry.split_dim is not None: + tensor_to_load = loaded_weight.chunk( + plan_entry.split_chunks, dim=plan_entry.split_dim + )[plan_entry.split_index] + + if self._fused_load_style == "per_expert": + _load_fused_expert_tensor( + param, + tensor_to_load, + shard_id=plan_entry.shard_id, + num_experts=self._num_experts, + ep_rank=self._ep_rank, + ep_size=self._ep_size, + ) + else: + if self._transpose_local_tensor_non_bias and "bias" not in mapped_name: + tensor_to_load = tensor_to_load.transpose(-2, -1) + + local_num_experts = param.shape[0] + assert local_num_experts * self._ep_size == tensor_to_load.shape[0] + local_experts = tensor_to_load[ + local_num_experts + * self._ep_rank : local_num_experts + * (self._ep_rank + 1) + ] + if tensor_to_load.dtype == torch.float8_e5m2: + default_weight_loader(param, local_experts.to(torch.bfloat16)) + else: + default_weight_loader(param, local_experts) + + loaded_any = True + + if not loaded_any: + assert mapped_name is not None + self._raise_unloaded_match(name, mapped_name) + return mapped_name + + def load(self, name: str, loaded_weight: torch.Tensor) -> str: + fused_mapped_name = self._load_fused(name, loaded_weight) + if fused_mapped_name is not None: + return fused_mapped_name + + expert_mapped_name = self._load_expert(name, loaded_weight) + if expert_mapped_name is not None: + return expert_mapped_name + + self._raise_unmatched(name) + + +def build_moe_checkpoint_loader( + *, + params_dict: dict[str, nn.Parameter], + expert_schema: ExpertCheckpointSchema | None = None, + fused_schema: ExpertCheckpointSchema | None = None, + num_experts: int | None = None, + ep_rank: int = 0, + ep_size: int = 1, + fused_gate_up_as_w13: bool = False, + include_bias: bool = False, + fused_load_style: str = "per_expert", + transpose_local_tensor_non_bias: bool = False, +) -> MoECheckpointLoader: + expert_plan: Sequence[ExpertWeightPlanEntry] = () + global_expert_plan: Sequence[CheckpointPlanEntry] = () + if expert_schema is not None: + if num_experts is None: + raise ValueError("num_experts is required when expert_schema is used") + expert_plan = _build_default_expert_plan( + expert_schema, + num_experts=num_experts, + ep_rank=ep_rank, + ep_size=ep_size, + ) + global_expert_plan = _build_global_expert_name_plan( + expert_schema, + num_experts=num_experts, + ) + + fused_plan: Sequence[FusedExpertWeightPlanEntry] = () + if fused_schema is not None: + fused_plan = _build_default_fused_plan( + fused_schema, + fused_gate_up_as_w13=fused_gate_up_as_w13, + include_bias=include_bias, + ) + + return MoECheckpointLoader( + params_dict=params_dict, + expert_plan=expert_plan, + global_expert_plan=global_expert_plan, + fused_plan=fused_plan, + num_experts=num_experts, + ep_rank=ep_rank, + ep_size=ep_size, + fused_load_style=fused_load_style, + transpose_local_tensor_non_bias=transpose_local_tensor_non_bias, + ) diff --git a/python/tokenspeed/runtime/layers/moe/schema.py b/python/tokenspeed/runtime/layers/moe/schema.py new file mode 100644 index 0000000..864eaa9 --- /dev/null +++ b/python/tokenspeed/runtime/layers/moe/schema.py @@ -0,0 +1,50 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from dataclasses import dataclass, field + + +@dataclass(frozen=True) +class ExpertCheckpointSchema: + gate_proj_name: str = "gate_proj" + up_proj_name: str = "up_proj" + down_proj_name: str = "down_proj" + gate_up_fused_name: str | None = None + extra_names: dict[str, str] = field(default_factory=dict) + + def get_semantic_name(self, semantic: str) -> str: + if semantic == "gate_proj": + return self.gate_proj_name + if semantic == "up_proj": + return self.up_proj_name + if semantic == "down_proj": + return self.down_proj_name + if semantic == "gate_up_fused": + if self.gate_up_fused_name is None: + raise KeyError("gate_up_fused_name is not defined") + return self.gate_up_fused_name + if semantic in self.extra_names: + return self.extra_names[semantic] + raise KeyError(f"Unknown MoE checkpoint semantic: {semantic}") + + def make_expert_weight_name(self, expert_id: int, semantic: str) -> str: + return f"experts.{expert_id}.{self.get_semantic_name(semantic)}." diff --git a/python/tokenspeed/runtime/layers/moe/topk.py b/python/tokenspeed/runtime/layers/moe/topk.py new file mode 100755 index 0000000..e0f99cf --- /dev/null +++ b/python/tokenspeed/runtime/layers/moe/topk.py @@ -0,0 +1,574 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from collections.abc import Callable +from dataclasses import dataclass +from enum import Enum, auto +from typing import Any, Literal, NamedTuple, Protocol, runtime_checkable + +import torch +import torch.nn.functional as F +from tokenspeed_kernel.thirdparty.cuda import routing_flash as cuda_routing_flash +from tokenspeed_kernel.thirdparty.triton import minimax_biased_grouped_topk + +from tokenspeed.runtime.moe.distribution_recorder import ( + get_global_expert_distribution_recorder, +) + + +class TopKOutputFormat(Enum): + STANDARD = auto() + BYPASSED = auto() + + def is_standard(self) -> bool: + return self == TopKOutputFormat.STANDARD + + def is_bypassed(self) -> bool: + return self == TopKOutputFormat.BYPASSED + + +@dataclass +class ExpertLocationDispatchInfo: + ep_dispatch_algorithm: Literal[ + "static", + "dynamic", + "fake", + "static_with_zero_expert", + "dynamic_with_zero_expert", + ] + # (num_logical_experts,) + partial_logical_to_rank_dispatch_physical_map: torch.Tensor | None + # (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: str | None = None, + expert_location_metadata: Any | None = 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: torch.Tensor | None, + info: ExpertLocationDispatchInfo | None, +): + if (info is not None) and (info.ep_dispatch_algorithm == "fake"): + router_logits = torch.randn_like(router_logits) + 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: ExpertLocationDispatchInfo | None, + num_experts: int | None = None, +) -> torch.Tensor: + if info is None: + return topk_ids + + if info.ep_dispatch_algorithm == "static": + return info.partial_logical_to_rank_dispatch_physical_map[topk_ids] + if info.ep_dispatch_algorithm == "static_with_zero_expert": + assert num_experts is not None + return _topk_ids_logical_to_physical_static_with_zero_expert( + topk_ids, info, num_experts + ) + if info.ep_dispatch_algorithm == "dynamic_with_zero_expert": + assert num_experts is not None + return _topk_ids_logical_to_physical_dynamic_with_zero_expert( + topk_ids, info, num_experts + ) + 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_with_zero_expert( + topk_ids: torch.Tensor, + info: ExpertLocationDispatchInfo, + num_experts: int, +) -> torch.Tensor: + topk_ids_original_shape = topk_ids.shape + topk_ids = topk_ids.flatten() + mask_less_than_num_experts = topk_ids < num_experts + converted_part = info.partial_logical_to_rank_dispatch_physical_map[ + topk_ids[mask_less_than_num_experts] + ] + topk_ids[mask_less_than_num_experts] = converted_part + return topk_ids.view(topk_ids_original_shape) + + +def _topk_ids_logical_to_physical_dynamic_with_zero_expert( + topk_ids: torch.Tensor, + info: ExpertLocationDispatchInfo, + num_experts: int, +) -> torch.Tensor: + topk_ids_original_shape = topk_ids.shape + device = topk_ids.device + topk_ids = topk_ids.flatten() + + mask_less_than_num_experts = topk_ids < num_experts + topk_ids_to_convert = topk_ids[mask_less_than_num_experts] + + chosen_dispatch_index = ( + torch.randint( + 0, 65536, topk_ids_to_convert.shape, dtype=torch.int32, device=device + ) + % info.partial_logical_to_all_physical_map_num_valid[topk_ids_to_convert] + ) + converted_topk_ids = info.partial_logical_to_all_physical_map[ + topk_ids_to_convert, chosen_dispatch_index + ] + topk_ids[mask_less_than_num_experts] = converted_topk_ids + return topk_ids.view(topk_ids_original_shape) + + +def _topk_ids_logical_to_physical_dynamic( + topk_ids: torch.Tensor, + info: 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] + return topk_ids.view(topk_ids_original_shape) + + +def _mask_topk_ids_padded_region( + topk_ids: torch.Tensor, + num_token_non_padded: torch.Tensor | None = 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 + + +def torch_native_fused_topk( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + correction_bias: torch.Tensor | None = 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=}" + 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 grouped_topk_gpu( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, + num_expert_group: int | None = None, + topk_group: int | None = None, + num_fused_shared_experts: int = 0, + routed_scaling_factor: float | None = None, +): + assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" + + scores = torch.softmax(gating_output, dim=-1) + num_token = scores.shape[0] + num_experts = scores.shape[1] + group_scores = scores.view(num_token, num_expert_group, -1).max(dim=-1).values + group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[1] + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, group_idx, 1) + score_mask = ( + group_mask.unsqueeze(-1) + .expand(num_token, num_expert_group, scores.shape[-1] // num_expert_group) + .reshape(num_token, -1) + ) + tmp_scores = scores.masked_fill(~score_mask.bool(), 0.0) + 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, + ) + factor = routed_scaling_factor or 1.0 + topk_weights[:, -1] = topk_weights[:, :-1].sum(dim=-1) / 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 routed_scaling_factor is not None: + topk_weights *= routed_scaling_factor + + return topk_weights.to(torch.float32), topk_ids.to(torch.int32) + + +@dataclass +class TopKConfig: + top_k: int + use_grouped_topk: bool = False + topk_group: int | None = None + num_expert_group: int | None = None + renormalize: bool = True + num_fused_shared_experts: int = 0 + custom_routing_function: Callable | None = None + correction_bias: torch.Tensor | None = None + torch_native: bool = False + routed_scaling_factor: float | None = None + output_format: TopKOutputFormat | None = None + zero_expert_num: int | None = 0 + topk_indices_dtype: torch.dtype | None = torch.int32 + + +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 BypassedTopKOutput(NamedTuple): + """Bypassed top-k output format.""" + + hidden_states: torch.Tensor + router_logits: torch.Tensor + topk_config: TopKConfig + num_token_non_padded: torch.Tensor | None = None + expert_location_dispatch_info: ExpertLocationDispatchInfo | None = None + + @property + def format(self) -> TopKOutputFormat: + return 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 TopK(torch.nn.Module): + + def __init__( + self, + top_k: int, + *, + use_grouped_topk: bool = False, + topk_group: int | None = None, + num_expert_group: int | None = None, + renormalize: bool = True, + num_fused_shared_experts: int = 0, + custom_routing_function: Callable | None = None, + correction_bias: torch.Tensor | None = None, + routed_scaling_factor: float | None = None, + output_format: TopKOutputFormat | None = None, + zero_expert_num: int | None = 0, + topk_indices_dtype=torch.int32, + ): + 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, + output_format=output_format, + zero_expert_num=zero_expert_num, + topk_indices_dtype=topk_indices_dtype, + ) + + def forward( + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + *, + num_token_non_padded: torch.Tensor | None = None, + expert_location_dispatch_info: ExpertLocationDispatchInfo | None = None, + ) -> TopKOutput: + if self.topk_config.output_format is not None: + output_format = self.topk_config.output_format + else: + output_format = TopKOutputFormat.STANDARD + + if output_format == TopKOutputFormat.BYPASSED: + 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 empty_topk_output( + self, + device: torch.device, + *, + hidden_states: torch.Tensor | None = None, + router_logits: torch.Tensor | None = None, + ) -> TopKOutput: + output_format = self.topk_config.output_format or TopKOutputFormat.STANDARD + if output_format.is_bypassed(): + if hidden_states is None: + hidden_states = torch.empty((0, 0), dtype=torch.float32, device=device) + if router_logits is None: + router_logits = torch.empty((0, 0), dtype=torch.float32, device=device) + return BypassedTopKOutput( + hidden_states=hidden_states, + router_logits=router_logits, + topk_config=self.topk_config, + ) + + 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=self.topk_config.topk_indices_dtype, + device=device, + ) + if router_logits is None: + router_logits = torch.empty((0, topk), dtype=torch.float32, device=device) + return StandardTopKOutput(topk_weights, topk_idx, router_logits) + + +def select_experts( + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + topk_config: TopKConfig, + *, + num_token_non_padded: torch.Tensor | None = None, + expert_location_dispatch_info: ExpertLocationDispatchInfo | None = 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 + + router_logits, correction_bias = 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_gpu( + hidden_states, + 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, + ) + else: + mapped_in_kernel = False + logical_to_physical_map = None + if ( + expert_location_dispatch_info is not None + and expert_location_dispatch_info.ep_dispatch_algorithm == "static" + ): + logical_to_physical_map = ( + expert_location_dispatch_info.partial_logical_to_rank_dispatch_physical_map + ) + mapped_in_kernel = True + topk_weights, topk_ids = minimax_biased_grouped_topk( + hidden_states, + router_logits, + 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, + logical_to_physical_map=logical_to_physical_map, + ) + if mapped_in_kernel: + expert_location_dispatch_info = None + + topk_ids = topk_ids_logical_to_physical( + topk_ids, + expert_location_dispatch_info, + num_experts=router_logits.shape[1], + ) + _mask_topk_ids_padded_region(topk_ids, num_token_non_padded) + 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 + topk_weights, topk_ids = torch_native_fused_topk( + hidden_states, + router_logits, + topk=top_k, + renormalize=renormalize, + correction_bias=correction_bias, + ) + if routed_scaling_factor is not None: + topk_weights *= routed_scaling_factor + elif correction_bias is not None: + # Bias-corrected top-k uses the CUDA fused_topk_bias kernel. + num_tokens = router_logits.shape[0] + topk_ids = torch.empty( + num_tokens, + top_k, + device=router_logits.device, + dtype=topk_config.topk_indices_dtype, + ) + topk_weights = torch.empty( + num_tokens, top_k, device=router_logits.device, dtype=torch.float32 + ) + num_real_experts = router_logits.shape[1] - topk_config.zero_expert_num + cuda_routing_flash( + router_logits, + correction_bias, + topk_ids, + topk_weights, + num_real_experts, + routed_scaling_factor, + renormalize, + ) + elif custom_routing_function is None: + topk_weights, topk_ids = torch_native_fused_topk( + hidden_states, + router_logits, + topk=top_k, + renormalize=renormalize, + ) + if routed_scaling_factor is not None: + topk_weights *= routed_scaling_factor + topk_ids = topk_ids_logical_to_physical( + topk_ids, + expert_location_dispatch_info, + num_experts=router_logits.shape[1], + ) + _mask_topk_ids_padded_region(topk_ids, num_token_non_padded) + + 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 + topk_weights, topk_ids = custom_routing_function( + hidden_states=hidden_states, + gating_output=router_logits, + topk=top_k, + renormalize=renormalize, + ) + if routed_scaling_factor is not None: + topk_weights *= routed_scaling_factor + + get_global_expert_distribution_recorder().on_select_experts(topk_ids=topk_ids) + + return StandardTopKOutput(topk_weights, topk_ids, router_logits) diff --git a/python/tokenspeed/runtime/layers/moe/types.py b/python/tokenspeed/runtime/layers/moe/types.py new file mode 100644 index 0000000..f00f6c2 --- /dev/null +++ b/python/tokenspeed/runtime/layers/moe/types.py @@ -0,0 +1,43 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class MoELayerSpec: + top_k: int + num_experts: int + num_local_experts: int + hidden_size: int + intermediate_size: int + activation: str + tp_rank: int + tp_size: int + ep_rank: int + ep_size: int + prefix: str = "" + a2a_backend: str = "none" + + @property + def use_deepep(self) -> bool: + return self.a2a_backend == "deepep" diff --git a/python/tokenspeed/runtime/layers/moe/utils.py b/python/tokenspeed/runtime/layers/moe/utils.py new file mode 100755 index 0000000..4dfdab3 --- /dev/null +++ b/python/tokenspeed/runtime/layers/moe/utils.py @@ -0,0 +1,168 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import logging +from enum import Enum, IntEnum +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from tokenspeed.runtime.utils.server_args import ServerArgs + +logger = logging.getLogger(__name__) + + +class RoutingMethodType(IntEnum): + Default = 0 + Renormalize = 1 + DeepSeekV3 = 2 + Llama4 = 3 + RenormalizeNaive = 4 + TopK = 5 + SigmoidRenorm = 6 + MiniMax2 = 7 + Unspecified = 8 + + +class All2AllBackend(Enum): + + NONE = "none" + DEEPEP = "deepep" + FLASHINFER_NVLINK_ONE_SIDED = "flashinfer_nvlink_one_sided" + + @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 == All2AllBackend.NONE + + def is_deepep(self): + return self == All2AllBackend.DEEPEP + + def is_flashinfer_nvlink_one_sided(self): + return self == All2AllBackend.FLASHINFER_NVLINK_ONE_SIDED + + +class MoeBackend(Enum): + + AUTO = "auto" + TRITON = "triton" + GLUON = "gluon" + FLASHINFER_TRTLLM = "flashinfer_trtllm" + FLASHINFER_CUTLASS = "flashinfer_cutlass" + FLASHINFER_CUTEDSL = "flashinfer_cutedsl" + + DEEP_GEMM_MEGA_MOE = "deep_gemm_mega_moe" + MEGA_MOE = "mega_moe" + + def is_auto(self): + return self == MoeBackend.AUTO + + def is_triton(self): + return self == MoeBackend.TRITON + + def is_gluon(self): + return self == MoeBackend.GLUON + + def is_flashinfer_trtllm(self): + return self == MoeBackend.FLASHINFER_TRTLLM + + def is_flashinfer_cutlass(self): + return self == MoeBackend.FLASHINFER_CUTLASS + + def is_flashinfer_cutedsl(self): + return self == MoeBackend.FLASHINFER_CUTEDSL + + def is_deep_gemm_mega_moe(self): + return self in (MoeBackend.DEEP_GEMM_MEGA_MOE, MoeBackend.MEGA_MOE) + + def is_mega_moe(self): + return self.is_deep_gemm_mega_moe() + + +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 + + +ALL2ALL_BACKEND: All2AllBackend | None = None +MOE_BACKEND: MoeBackend | None = None +DISABLE_FLASHINFER_CUTLASS_MOE_FP4_ALLGATHER: bool | None = None + + +def initialize_moe_config(server_args: ServerArgs): + global ALL2ALL_BACKEND + global MOE_BACKEND + global DISABLE_FLASHINFER_CUTLASS_MOE_FP4_ALLGATHER + + ALL2ALL_BACKEND = All2AllBackend(server_args.all2all_backend) + MOE_BACKEND = MoeBackend(server_args.moe_backend) + DISABLE_FLASHINFER_CUTLASS_MOE_FP4_ALLGATHER = ( + server_args.disable_flashinfer_cutlass_moe_fp4_allgather + ) + + +def get_all2all_backend() -> All2AllBackend: + global ALL2ALL_BACKEND + if ALL2ALL_BACKEND is None: + logger.warning("ALL2ALL_BACKEND is not initialized, using default backend") + ALL2ALL_BACKEND = All2AllBackend.NONE + return ALL2ALL_BACKEND + + +def get_moe_backend() -> MoeBackend: + global MOE_BACKEND + if MOE_BACKEND is None: + logger.warning("MOE_BACKEND is not initialized, using auto backend") + MOE_BACKEND = MoeBackend.AUTO + return MOE_BACKEND diff --git a/python/tokenspeed/runtime/layers/moe/weights/__init__.py b/python/tokenspeed/runtime/layers/moe/weights/__init__.py new file mode 100644 index 0000000..1df40a8 --- /dev/null +++ b/python/tokenspeed/runtime/layers/moe/weights/__init__.py @@ -0,0 +1,94 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import torch + +from tokenspeed.runtime.layers.moe.weights.fp8 import create_fp8_block_scale_inverses +from tokenspeed.runtime.layers.moe.weights.mxfp4 import ( + create_mxfp4_fp8_input_scales, + create_mxfp4_weight_pair, +) +from tokenspeed.runtime.layers.moe.weights.mxint4 import create_mxint4_weight_pair +from tokenspeed.runtime.layers.moe.weights.nvfp4 import create_nvfp4_weight_pair +from tokenspeed.runtime.layers.moe.weights.unquant import create_dense_weight_pair + + +def create_layer_weights( + spec, + layer, + quant_kind: str, + quant_config, + *, + with_bias: bool = False, + solution: str | None = None, +) -> None: + if quant_kind == "unquant": + create_dense_weight_pair( + spec, + layer, + params_dtype=torch.get_default_dtype(), + with_bias=with_bias, + ) + return + + if quant_kind == "fp8": + ispp = create_dense_weight_pair( + spec, + layer, + params_dtype=torch.float8_e4m3fn, + with_bias=with_bias, + ) + create_fp8_block_scale_inverses( + spec, + layer, + intermediate_size_per_partition=ispp, + block_shape=quant_config.weight_block_size, + ) + return + + if quant_kind == "nvfp4": + create_nvfp4_weight_pair( + spec, + layer, + group_size=quant_config.group_size, + ) + return + + if quant_kind == "mxfp4": + create_mxfp4_weight_pair(spec, layer, with_bias=with_bias, solution=solution) + if quant_config.is_w4a8_fp8: + create_mxfp4_fp8_input_scales(layer, spec.num_local_experts) + return + + if quant_kind == "mxint4": + weight_quant = quant_config.target_scheme_map["Linear"]["weights"] + create_mxint4_weight_pair( + spec, + layer, + group_size=weight_quant.group_size, + ) + return + + raise RuntimeError(f"Unsupported MoE quant kind: {quant_kind}") + + +__all__ = [ + "create_layer_weights", +] diff --git a/python/tokenspeed/runtime/layers/moe/weights/fp8.py b/python/tokenspeed/runtime/layers/moe/weights/fp8.py new file mode 100644 index 0000000..7782c0c --- /dev/null +++ b/python/tokenspeed/runtime/layers/moe/weights/fp8.py @@ -0,0 +1,68 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import torch +from torch import nn + +from tokenspeed.runtime.layers.moe.types import MoELayerSpec +from tokenspeed.runtime.layers.moe.weights.loaders import ( + make_weight_loader, + round_up, +) +from tokenspeed.runtime.utils import set_weight_attrs + + +def create_fp8_block_scale_inverses( + spec: MoELayerSpec, + layer: nn.Module, + *, + intermediate_size_per_partition: int, + block_shape: tuple[int, int], +) -> None: + block_n, block_k = block_shape + w13_weight_scale = torch.nn.Parameter( + torch.ones( + spec.num_local_experts, + 2 * round_up(intermediate_size_per_partition, block_n) // block_n, + round_up(spec.hidden_size, block_k) // block_k, + dtype=torch.float32, + ), + requires_grad=False, + ) + w2_weight_scale = torch.nn.Parameter( + torch.ones( + spec.num_local_experts, + round_up(spec.hidden_size, block_n) // block_n, + round_up(intermediate_size_per_partition, block_k) // 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) + + weight_loader = make_weight_loader(spec) + set_weight_attrs(w13_weight_scale, {"weight_loader": weight_loader}) + set_weight_attrs(w2_weight_scale, {"weight_loader": weight_loader}) + + +__all__ = ["create_fp8_block_scale_inverses"] diff --git a/python/tokenspeed/runtime/layers/moe/weights/loaders.py b/python/tokenspeed/runtime/layers/moe/weights/loaders.py new file mode 100644 index 0000000..840c4bb --- /dev/null +++ b/python/tokenspeed/runtime/layers/moe/weights/loaders.py @@ -0,0 +1,282 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from collections.abc import Callable +from functools import partial + +import torch + +from tokenspeed.runtime.layers.moe.types import MoELayerSpec + + +def preserve_e8m0_bytes_for_uint8_param( + dst: torch.Tensor, + src: torch.Tensor, +) -> torch.Tensor: + e8m0_dtype = getattr(torch, "float8_e8m0fnu", None) + if e8m0_dtype is not None and dst.dtype == torch.uint8 and src.dtype == e8m0_dtype: + return src.view(torch.uint8) + return src + + +def load_w13( + expert_data: torch.Tensor, + loaded_weight: torch.Tensor, + shard_id: str, + shard_dim: int, + tp_rank: int, + is_bias: bool, + use_presharded_weights: bool, + do_transpose: bool, + tp_size: int = 1, + load_up_proj_weight_first: bool = False, +) -> None: + if shard_id not in {"w1", "w3", "w13"}: + raise ValueError(f"Unexpected w13 shard_id: {shard_id}") + + if is_bias: + shard_dim = -1 + + if shard_id in {"w1", "w3"}: + shard_size = expert_data.shape[shard_dim] // 2 + else: + shard_size = expert_data.shape[shard_dim] + + switch_w13 = load_up_proj_weight_first + if (switch_w13 and shard_id == "w1") or (not switch_w13 and shard_id == "w3"): + start = shard_size + else: + start = 0 + + if not use_presharded_weights: + if not is_bias and do_transpose: + loaded_weight = loaded_weight.transpose(-2, -1) + if tp_size > 1: + # Derive the unpadded shard size from the checkpoint tensor. + # expert_data may be Blackwell-padded; checkpoint is not. + load_shard = loaded_weight.shape[shard_dim] // tp_size + loaded_weight = loaded_weight.narrow( + shard_dim, load_shard * tp_rank, load_shard + ) + else: + loaded_weight = loaded_weight.narrow( + shard_dim, shard_size * tp_rank, shard_size + ) + + expert_data = expert_data.narrow(shard_dim, start, shard_size) + dst = expert_data.narrow(shard_dim, 0, loaded_weight.shape[shard_dim]) + loaded_weight = preserve_e8m0_bytes_for_uint8_param(dst, loaded_weight) + dst.copy_(loaded_weight) + + +def load_w2( + expert_data: torch.Tensor, + loaded_weight: torch.Tensor, + shard_id: str, + shard_dim: int, + tp_rank: int, + is_bias: bool, + use_presharded_weights: bool, + do_transpose: bool, + tp_size: int = 1, +) -> None: + 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 shard_id != "w2": + raise ValueError(f"shard_id must be 'w2', got {shard_id}") + + if is_bias: + shard_dim = -1 + shard_size = expert_data.shape[-1] + else: + shard_size = expert_data.shape[shard_dim] + + if not use_presharded_weights: + if not is_bias and do_transpose: + loaded_weight = loaded_weight.transpose(-2, -1) + if is_bias: + load_shard = shard_size + elif tp_size > 1: + load_shard = loaded_weight.shape[shard_dim] // tp_size + else: + load_shard = shard_size + start = 0 if is_bias else load_shard * tp_rank + loaded_weight = loaded_weight.narrow(shard_dim, start, load_shard) + + dst = expert_data.narrow(shard_dim, 0, loaded_weight.shape[shard_dim]) + loaded_weight = preserve_e8m0_bytes_for_uint8_param(dst, loaded_weight) + dst.copy_(loaded_weight) + + +def get_shard_dim(param: torch.Tensor, shard_id: str, do_transpose: bool) -> int: + is_transposed = getattr(param, "is_transposed", False) + if do_transpose: + is_transposed = True + + shard_dim = {"w1": 0, "w2": 1, "w3": 0, "w13": 0}[shard_id] + if is_transposed: + shard_dim = int(not shard_dim) + return shard_dim + + +def load_model_weight( + param: torch.Tensor, + loaded_weight: torch.Tensor, + shard_id: str, + local_expert_id: int, + tp_rank: int, + is_bias: bool, + use_presharded_weights: bool, + do_transpose: bool, + tp_size: int = 1, +) -> None: + expert_data = param.data[local_expert_id] + shard_dim = get_shard_dim(param, shard_id, do_transpose) + if shard_id == "w2": + load_w2( + expert_data, + loaded_weight, + shard_id, + shard_dim, + tp_rank, + is_bias, + use_presharded_weights, + do_transpose, + tp_size=tp_size, + ) + elif shard_id in {"w1", "w3", "w13"}: + load_w13( + expert_data, + loaded_weight, + shard_id, + shard_dim, + tp_rank, + is_bias, + use_presharded_weights, + do_transpose, + tp_size=tp_size, + ) + else: + raise ValueError(f"Unknown shard_id: {shard_id}") + + +def load_group_weight_scale( + param: torch.Tensor, + loaded_weight: torch.Tensor, + local_expert_id: int, + shard_id: str, + tp_rank: int, + do_transpose: bool, + tp_size: int = 1, +) -> None: + load_model_weight( + param, + loaded_weight, + shard_id, + local_expert_id, + tp_rank, + False, + False, + do_transpose, + tp_size=tp_size, + ) + + +def load_per_tensor_weight_scale( + param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + shard_id: str, + local_expert_id: int, +) -> None: + if shard_id in {"w1", "w3"}: + idx = 0 if shard_id == "w1" else 1 + param.data[local_expert_id][idx] = loaded_weight + elif shard_id == "w2": + param.data[local_expert_id] = loaded_weight + else: + raise ValueError(f"Unknown shard_id: {shard_id}") + + +def load_per_tensor_input_scale( + param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + shard_id: str, + local_expert_id: int, +) -> None: + value = loaded_weight.detach().to(torch.float32).reshape(()) + if shard_id in {"w1", "w3"}: + prev = param.data[local_expert_id] + param.data[local_expert_id] = torch.maximum(prev, value) + elif shard_id == "w2": + param.data[local_expert_id] = value + else: + raise ValueError(f"Unknown shard_id for input scale: {shard_id}") + + +def make_weight_loader( + spec: MoELayerSpec, + *, + is_bias: bool = False, + do_transpose: bool = False, + use_presharded_weights: bool = False, +) -> Callable: + return partial( + load_model_weight, + tp_rank=spec.tp_rank, + is_bias=is_bias, + use_presharded_weights=use_presharded_weights, + do_transpose=do_transpose, + tp_size=spec.tp_size, + ) + + +def make_group_scale_loader( + spec: MoELayerSpec, + *, + do_transpose: bool = False, +) -> Callable: + return partial( + load_group_weight_scale, + tp_rank=spec.tp_rank, + do_transpose=do_transpose, + tp_size=spec.tp_size, + ) + + +def per_tensor_scale_loader() -> Callable: + return load_per_tensor_weight_scale + + +def round_up(value: int, multiple: int) -> int: + return (value + multiple - 1) // multiple * multiple + + +__all__ = [ + "load_per_tensor_input_scale", + "make_group_scale_loader", + "make_weight_loader", + "per_tensor_scale_loader", + "round_up", +] diff --git a/python/tokenspeed/runtime/layers/moe/weights/mxfp4.py b/python/tokenspeed/runtime/layers/moe/weights/mxfp4.py new file mode 100644 index 0000000..6cb2aa6 --- /dev/null +++ b/python/tokenspeed/runtime/layers/moe/weights/mxfp4.py @@ -0,0 +1,139 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import torch +from tokenspeed_kernel.platform import current_platform +from torch import nn + +from tokenspeed.runtime.layers.moe.types import MoELayerSpec +from tokenspeed.runtime.layers.moe.weights.loaders import ( + load_per_tensor_input_scale, + make_weight_loader, + round_up, +) +from tokenspeed.runtime.utils import set_weight_attrs + +MXFP4_BLOCK = 32 + + +def create_mxfp4_weight_pair( + spec: MoELayerSpec, + layer: nn.Module, + *, + with_bias: bool = False, + solution: str | None = None, +) -> None: + ispp = spec.intermediate_size // spec.tp_size + platform = current_platform() + hidden_size_padded = spec.hidden_size + if platform.is_blackwell and solution == "flashinfer_trtllm": + ispp_padded = round_up(ispp, 256) + hidden_size_padded = round_up(spec.hidden_size, 256) + else: + ispp_padded = ( + round_up(ispp, 64) if platform.is_blackwell else round_up(ispp, MXFP4_BLOCK) + ) + + w13_weight = torch.nn.Parameter( + torch.zeros( + spec.num_local_experts, + 2 * ispp_padded, + hidden_size_padded // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ) + w13_weight_scale = torch.nn.Parameter( + torch.zeros( + spec.num_local_experts, + 2 * ispp_padded, + hidden_size_padded // MXFP4_BLOCK, + dtype=torch.uint8, + ), + requires_grad=False, + ) + w2_weight = torch.nn.Parameter( + torch.zeros( + spec.num_local_experts, + hidden_size_padded, + ispp_padded // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ) + w2_weight_scale = torch.nn.Parameter( + torch.zeros( + spec.num_local_experts, + hidden_size_padded, + ispp_padded // MXFP4_BLOCK, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + layer.register_parameter("w2_weight", w2_weight) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + + weight_loader = make_weight_loader(spec) + set_weight_attrs(w13_weight, {"weight_loader": weight_loader}) + set_weight_attrs(w13_weight_scale, {"weight_loader": weight_loader}) + set_weight_attrs(w2_weight, {"weight_loader": weight_loader}) + set_weight_attrs(w2_weight_scale, {"weight_loader": weight_loader}) + + if with_bias: + w13_weight_bias = torch.nn.Parameter( + torch.zeros(spec.num_local_experts, 2 * ispp_padded, dtype=torch.bfloat16), + requires_grad=False, + ) + w2_weight_bias = torch.nn.Parameter( + torch.zeros( + spec.num_local_experts, hidden_size_padded, dtype=torch.bfloat16 + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_bias", w13_weight_bias) + layer.register_parameter("w2_weight_bias", w2_weight_bias) + bias_loader = make_weight_loader(spec, is_bias=True) + set_weight_attrs(w13_weight_bias, {"weight_loader": bias_loader}) + set_weight_attrs(w2_weight_bias, {"weight_loader": bias_loader}) + + +def create_mxfp4_fp8_input_scales( + layer: nn.Module, + num_local_experts: int, +) -> None: + w13_input_scale = nn.Parameter( + torch.zeros(num_local_experts, dtype=torch.float32), + requires_grad=False, + ) + w2_input_scale = nn.Parameter( + torch.zeros(num_local_experts, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w13_input_scale", w13_input_scale) + layer.register_parameter("w2_input_scale", w2_input_scale) + set_weight_attrs(w13_input_scale, {"weight_loader": load_per_tensor_input_scale}) + set_weight_attrs(w2_input_scale, {"weight_loader": load_per_tensor_input_scale}) + + +__all__ = ["create_mxfp4_fp8_input_scales", "create_mxfp4_weight_pair"] diff --git a/python/tokenspeed/runtime/layers/moe/weights/mxint4.py b/python/tokenspeed/runtime/layers/moe/weights/mxint4.py new file mode 100644 index 0000000..f315c10 --- /dev/null +++ b/python/tokenspeed/runtime/layers/moe/weights/mxint4.py @@ -0,0 +1,124 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import torch +from torch import nn + +from tokenspeed.runtime.layers.moe.types import MoELayerSpec +from tokenspeed.runtime.layers.moe.weights.loaders import ( + make_group_scale_loader, + make_weight_loader, +) +from tokenspeed.runtime.utils import set_weight_attrs + +# Eight INT4 values per int32 word (32 // 4); the kernel and its repack are +# INT4-only, so this is a constant, not a config knob. +_PACKED_FACTOR = 8 + + +def _ignore_weight_loader(param, loaded_weight, **kwargs) -> None: + """No-op loader for checkpoint metadata tensors the kernel does not use.""" + del param, loaded_weight, kwargs + + +def create_mxint4_weight_pair( + spec: MoELayerSpec, + layer: nn.Module, + *, + group_size: int, +) -> None: + """Register per-expert INT4 pack-quantized weights with bf16 group scales. + + Tensors keep the natural ``[out, in // _PACKED_FACTOR]`` checkpoint layout + (gate/up fused along the output dim), so the shared MoE checkpoint loaders + fill them without transposition; the trtllm process-weights kernel rewrites + them into the kernel layout afterwards. Eight INT4 values are packed per + ``int32`` word, while scales hold one ``bfloat16`` value per ``group_size`` + input elements. + """ + ispp = spec.intermediate_size // spec.tp_size + + # Fused gate_up_proj (column parallel): [2 * intermediate, hidden // pack]. + w13_weight_packed = torch.nn.Parameter( + torch.empty( + spec.num_local_experts, + 2 * ispp, + spec.hidden_size // _PACKED_FACTOR, + dtype=torch.int32, + ), + requires_grad=False, + ) + # down_proj (row parallel): [hidden, intermediate // pack]. + w2_weight_packed = torch.nn.Parameter( + torch.empty( + spec.num_local_experts, + spec.hidden_size, + ispp // _PACKED_FACTOR, + dtype=torch.int32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_packed", w13_weight_packed) + layer.register_parameter("w2_weight_packed", w2_weight_packed) + + # Per-group bf16 scales: one scale per ``group_size`` input elements. + w13_weight_scale = torch.nn.Parameter( + torch.ones( + spec.num_local_experts, + 2 * ispp, + spec.hidden_size // group_size, + dtype=torch.bfloat16, + ), + requires_grad=False, + ) + w2_weight_scale = torch.nn.Parameter( + torch.ones( + spec.num_local_experts, + spec.hidden_size, + ispp // group_size, + dtype=torch.bfloat16, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + + weight_loader = make_weight_loader(spec) + scale_loader = make_group_scale_loader(spec) + set_weight_attrs(w13_weight_packed, {"weight_loader": weight_loader}) + set_weight_attrs(w2_weight_packed, {"weight_loader": weight_loader}) + set_weight_attrs(w13_weight_scale, {"weight_loader": scale_loader}) + set_weight_attrs(w2_weight_scale, {"weight_loader": scale_loader}) + + # compressed-tensors ships a per-proj ``weight_shape`` metadata tensor the + # kernel ignores. Register absorbers so the loader has a target (it raises + # on a matched expert tensor with no parameter); dropped in process_weights. + for shape_name in ("w13_weight_shape", "w2_weight_shape"): + shape_param = torch.nn.Parameter( + torch.empty(spec.num_local_experts, 2, dtype=torch.int32), + requires_grad=False, + ) + layer.register_parameter(shape_name, shape_param) + set_weight_attrs(shape_param, {"weight_loader": _ignore_weight_loader}) + + +__all__ = ["create_mxint4_weight_pair"] diff --git a/python/tokenspeed/runtime/layers/moe/weights/nvfp4.py b/python/tokenspeed/runtime/layers/moe/weights/nvfp4.py new file mode 100644 index 0000000..1ce94fa --- /dev/null +++ b/python/tokenspeed/runtime/layers/moe/weights/nvfp4.py @@ -0,0 +1,118 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import torch +from torch import nn + +from tokenspeed.runtime.layers.moe.types import MoELayerSpec +from tokenspeed.runtime.layers.moe.weights.loaders import ( + make_group_scale_loader, + make_weight_loader, + per_tensor_scale_loader, +) +from tokenspeed.runtime.utils import set_weight_attrs + + +def create_nvfp4_weight_pair( + spec: MoELayerSpec, + layer: nn.Module, + *, + group_size: int, +) -> None: + ispp = spec.intermediate_size // spec.tp_size + w13_weight = torch.nn.Parameter( + torch.empty( + spec.num_local_experts, + 2 * ispp, + spec.hidden_size // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ) + w2_weight = torch.nn.Parameter( + torch.empty( + spec.num_local_experts, + spec.hidden_size, + ispp // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + layer.register_parameter("w2_weight", w2_weight) + + w13_weight_scale = torch.nn.Parameter( + torch.empty( + spec.num_local_experts, + 2 * ispp, + spec.hidden_size // group_size, + dtype=torch.float8_e4m3fn, + ), + requires_grad=False, + ) + w2_weight_scale = torch.nn.Parameter( + torch.empty( + spec.num_local_experts, + spec.hidden_size, + ispp // group_size, + dtype=torch.float8_e4m3fn, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale", w13_weight_scale) + layer.register_parameter("w2_weight_scale", w2_weight_scale) + + w13_weight_scale_2 = torch.nn.Parameter( + torch.empty(spec.num_local_experts, 2, dtype=torch.float32), + requires_grad=False, + ) + w2_weight_scale_2 = torch.nn.Parameter( + torch.empty(spec.num_local_experts, dtype=torch.float32), + requires_grad=False, + ) + w13_input_scale = torch.nn.Parameter( + torch.empty(spec.num_local_experts, 2, dtype=torch.float32), + requires_grad=False, + ) + w2_input_scale = torch.nn.Parameter( + torch.empty(spec.num_local_experts, dtype=torch.float32), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale_2", w13_weight_scale_2) + layer.register_parameter("w2_weight_scale_2", w2_weight_scale_2) + layer.register_parameter("w13_input_scale", w13_input_scale) + layer.register_parameter("w2_input_scale", w2_input_scale) + + weight_loader = make_weight_loader(spec) + scale_loader = make_group_scale_loader(spec) + per_tensor_loader = per_tensor_scale_loader() + set_weight_attrs(w13_weight, {"weight_loader": weight_loader}) + set_weight_attrs(w2_weight, {"weight_loader": weight_loader}) + set_weight_attrs(w13_weight_scale, {"weight_loader": scale_loader}) + set_weight_attrs(w2_weight_scale, {"weight_loader": scale_loader}) + set_weight_attrs(w13_weight_scale_2, {"weight_loader": per_tensor_loader}) + set_weight_attrs(w2_weight_scale_2, {"weight_loader": per_tensor_loader}) + set_weight_attrs(w13_input_scale, {"weight_loader": per_tensor_loader}) + set_weight_attrs(w2_input_scale, {"weight_loader": per_tensor_loader}) + + +__all__ = ["create_nvfp4_weight_pair"] diff --git a/python/tokenspeed/runtime/layers/moe/weights/unquant.py b/python/tokenspeed/runtime/layers/moe/weights/unquant.py new file mode 100644 index 0000000..f547ee8 --- /dev/null +++ b/python/tokenspeed/runtime/layers/moe/weights/unquant.py @@ -0,0 +1,82 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import torch +from torch import nn + +from tokenspeed.runtime.layers.moe.types import MoELayerSpec +from tokenspeed.runtime.layers.moe.weights.loaders import make_weight_loader +from tokenspeed.runtime.utils import set_weight_attrs + + +def create_dense_weight_pair( + spec: MoELayerSpec, + layer: nn.Module, + *, + params_dtype: torch.dtype, + with_bias: bool = False, +) -> int: + ispp = spec.intermediate_size // spec.tp_size + w13_weight = torch.nn.Parameter( + torch.empty( + spec.num_local_experts, + 2 * ispp, + spec.hidden_size, + dtype=params_dtype, + ), + requires_grad=False, + ) + w2_weight = torch.nn.Parameter( + torch.empty( + spec.num_local_experts, + spec.hidden_size, + ispp, + dtype=params_dtype, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight", w13_weight) + layer.register_parameter("w2_weight", w2_weight) + + weight_loader = make_weight_loader(spec) + set_weight_attrs(w13_weight, {"weight_loader": weight_loader}) + set_weight_attrs(w2_weight, {"weight_loader": weight_loader}) + + if with_bias: + w13_weight_bias = torch.nn.Parameter( + torch.zeros(spec.num_local_experts, 2 * ispp, dtype=params_dtype), + requires_grad=False, + ) + w2_weight_bias = torch.nn.Parameter( + torch.zeros(spec.num_local_experts, spec.hidden_size, dtype=params_dtype), + requires_grad=False, + ) + layer.register_parameter("w13_weight_bias", w13_weight_bias) + layer.register_parameter("w2_weight_bias", w2_weight_bias) + bias_loader = make_weight_loader(spec, is_bias=True) + set_weight_attrs(w13_weight_bias, {"weight_loader": bias_loader}) + set_weight_attrs(w2_weight_bias, {"weight_loader": bias_loader}) + + return ispp + + +__all__ = ["create_dense_weight_pair"] diff --git a/python/tokenspeed/runtime/layers/paged_attention.py b/python/tokenspeed/runtime/layers/paged_attention.py new file mode 100755 index 0000000..4be47f1 --- /dev/null +++ b/python/tokenspeed/runtime/layers/paged_attention.py @@ -0,0 +1,129 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Paged attention.""" + +from collections.abc import Sequence + +import torch +from torch import nn + +from tokenspeed.runtime.execution.context import ForwardContext + + +class PagedAttention(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, + group_id: 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 + # Flat KV-cache group ("" -> single-table fallback in the backend). + # TODO(radix-removal): make group_id mandatory once flat is the only path. + self.group_id = group_id + self.k_scale = None + self.v_scale = None + + def forward( + self, + q, + k, + v, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + save_kv_cache: bool = True, + **kwargs, + ): + if k is not None: + # For cross-layer sharing, kv can be None + if v is None: + raise ValueError("v must be provided when k is provided.") + if "k_pe" 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) + v = v.view(-1, self.tp_v_head_num, self.v_head_dim) + + return ctx.attn_backend.forward( + q, + k, + v, + self, + out_cache_loc, + ctx.token_to_kv_pool, + ctx.forward_mode, + ctx.bs, + save_kv_cache, + **kwargs, + ) + + +def validate_paged_cache_group_ids( + model: nn.Module, + paged_cache_group_specs: Sequence, +) -> None: + """Fail fast (ValueError) when a pool publishing more than one paged-cache + group meets a PagedAttention layer whose group_id is empty or unknown -- + instead of a KeyError deep in the backend, possibly during graph capture. + """ + group_ids = {str(spec.group_id) for spec in paged_cache_group_specs} + if len(group_ids) <= 1: + return + model_name = type(model).__name__ + for name, module in model.named_modules(): + if not isinstance(module, PagedAttention): + continue + if not module.group_id: + raise ValueError( + f"{model_name}: attention layer {name!r} (layer_id=" + f"{module.layer_id}) has empty group_id but the KV pool " + f"publishes {len(group_ids)} paged-cache groups " + f"{sorted(group_ids)}; pass group_id= to " + "PagedAttention (see gpt_oss.py)." + ) + if module.group_id not in group_ids: + raise ValueError( + f"{model_name}: attention layer {name!r} (layer_id=" + f"{module.layer_id}) has group_id={module.group_id!r} which " + "is not among the KV pool's paged-cache groups " + f"{sorted(group_ids)}." + ) diff --git a/python/tokenspeed/runtime/layers/parameter.py b/python/tokenspeed/runtime/layers/parameter.py new file mode 100755 index 0000000..32422f3 --- /dev/null +++ b/python/tokenspeed/runtime/layers/parameter.py @@ -0,0 +1,475 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Parameter helpers for sharded and packed layer weights.""" + +from collections.abc import Callable +from fractions import Fraction + +import torch +from torch.nn import Parameter + +from tokenspeed.runtime.utils import get_colorful_logger + +__all__ = [ + "BaseWeightParameter", + "PackedWeightParameter", + "PerTensorScaleParameter", + "ModelWeightParameter", + "ChannelQuantScaleParameter", + "GroupQuantScaleParameter", + "PackedColumnParameter", + "RowParallelWeightParameter", +] + +logger = get_colorful_logger(__name__) + + +def _check_shape_match(actual: torch.Size, expected: torch.Size) -> None: + if actual != expected: + raise ValueError(f"Shape mismatch: {actual} != {expected}.") + + +class BaseWeightParameter(Parameter): + """Base parameter for TokenSpeed linear layers with custom weight loading.""" + + 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 parameter wrapper with a weight-loader callback.""" + self._weight_loader = weight_loader + + @property + def weight_loader(self) -> Callable: + return self._weight_loader + + def _assert_and_load(self, loaded_weight: torch.Tensor) -> None: + _check_shape_match(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 _ColumnParallelWeightParameter(BaseWeightParameter): + """Shared column-parallel weight-loading helpers.""" + + def __init__(self, output_dim: int, **kwargs): + self._output_dim = output_dim + super().__init__(**kwargs) + + @property + def output_dim(self) -> int: + 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] + loaded_weight = loaded_weight.narrow( + self.output_dim, tp_rank * shard_size, shard_size + ) + _check_shape_match(self.data.shape, loaded_weight.shape) + self.data.copy_(loaded_weight) + + def load_merged_column_weight( + self, loaded_weight: torch.Tensor, tp_rank: int, **kwargs + ): + + shard_offset = kwargs.get("shard_offset") + shard_size = kwargs.get("shard_size") + use_presharded_weights = kwargs.get("use_presharded_weights") + if ( + isinstance(self, (PackedColumnParameter, PackedWeightParameter)) + 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) + if not use_presharded_weights: + loaded_weight = loaded_weight.narrow( + self.output_dim, tp_rank * shard_size, shard_size + ) + _check_shape_match(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, PackedWeightParameter)) + 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 not use_presharded_weights: + loaded_weight = loaded_weight.narrow( + self.output_dim, shard_id * shard_size, shard_size + ) + + _check_shape_match(param_data.shape, loaded_weight.shape) + param_data.copy_(loaded_weight) + + +class RowParallelWeightParameter(BaseWeightParameter): + """Parameter class with row-parallel weight-loading support.""" + + def __init__(self, input_dim: int, **kwargs): + self._input_dim = input_dim + super().__init__(**kwargs) + + @property + def input_dim(self) -> int: + 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] + 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) + + _check_shape_match(self.data.shape, loaded_weight.shape) + self.data.copy_(loaded_weight) + + +class ModelWeightParameter(_ColumnParallelWeightParameter, RowParallelWeightParameter): + """ + Parameter class for linear layer weights. Uses both column and + row parallelism. + """ + + pass + + +class GroupQuantScaleParameter( + _ColumnParallelWeightParameter, RowParallelWeightParameter +): + """ + Parameter class for weight scales loaded for weights with + grouped quantization. Uses both column and row parallelism. + """ + + pass + + +class ChannelQuantScaleParameter(_ColumnParallelWeightParameter): + """ + Parameter class for weight scales loaded for weights with + channel-wise quantization. Equivalent to _ColumnParallelWeightParameter. + """ + + pass + + +class PerTensorScaleParameter(BaseWeightParameter): + """ + 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: str | int) -> int: + if isinstance(shard_id, int): + return shard_id + + # if not int, assume shard_id for qkv + # map to int and return + if not isinstance(shard_id, str): + raise TypeError( + f"shard_id must be a str or int, got {type(shard_id).__name__}." + ) + if shard_id not in self.qkv_idxs: + raise ValueError(f"Invalid qkv shard_id: {shard_id}.") + 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: 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: + if loaded_weight.shape[0] != 1: + raise ValueError( + f"Expected scale shard with first dimension 1, got {loaded_weight.shape}." + ) + loaded_weight = loaded_weight[0] + + param_data = param_data[shard_id] + _check_shape_match(param_data.shape, loaded_weight.shape) + param_data.copy_(loaded_weight) + + +class PackedColumnParameter(_ColumnParallelWeightParameter): + """ + Parameter for model parameters which are packed on disk + and support column parallelism only. See PackedWeightParameter + for more details on the packed properties. + """ + + def __init__( + self, + packed_factor: int | Fraction, + packed_dim: int, + marlin_tile_size: int | None = 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 PackedWeightParameter(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: int | Fraction, + packed_dim: int, + marlin_tile_size: int | None = 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: BaseWeightParameter, input_dim: int, output_dim: int, **kwargs +) -> BaseWeightParameter: + """ + 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: + if param.data.dim() != 2: + raise ValueError( + "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: + if curr_output_dim is None: + raise ValueError("either input or output dim must be set") + curr_input_dim = (curr_output_dim + 1) % 2 + if curr_output_dim is None: + if curr_input_dim is None: + raise ValueError("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: + if not ( + hasattr(param, "packed_dim") + and param.packed_dim == perm[kwargs["packed_dim"]] + ): + raise ValueError( + "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 + + +class BlockQuantScaleParameter( + _ColumnParallelWeightParameter, RowParallelWeightParameter +): + """ + Parameter class for weight scales loaded for weights with + block-wise quantization. Uses both column and row parallelism. + """ + + pass diff --git a/python/tokenspeed/runtime/layers/quantization/__init__.py b/python/tokenspeed/runtime/layers/quantization/__init__.py new file mode 100644 index 0000000..1a0d95c --- /dev/null +++ b/python/tokenspeed/runtime/layers/quantization/__init__.py @@ -0,0 +1,60 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from enum import Enum + +from tokenspeed.runtime.layers.quantization.base_config import ( # noqa: F401 + QuantizationConfig, + QuantizeMethodBase, +) +from tokenspeed.runtime.layers.quantization.compressed_tensors.compressed_tensors import ( + CompressedTensorsConfig, +) +from tokenspeed.runtime.layers.quantization.fp8 import Fp8Config +from tokenspeed.runtime.layers.quantization.mxfp4 import Mxfp4Config +from tokenspeed.runtime.layers.quantization.nvfp4 import Nvfp4Config +from tokenspeed.runtime.layers.quantization.w8a8_fp8 import W8A8Fp8Config + +BASE_QUANTIZATION_METHODS: dict[str, type[QuantizationConfig]] = { + "fp8": Fp8Config, + "w8a8_fp8": W8A8Fp8Config, + "compressed-tensors": CompressedTensorsConfig, + "nvfp4": Nvfp4Config, + "mxfp4": Mxfp4Config, +} + + +QUANTIZATION_METHODS = BASE_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())}" + ) + return QUANTIZATION_METHODS[quantization] + + +class FusedMoeWeightScaleSupported(Enum): + TENSOR = "tensor" + CHANNEL = "channel" + GROUP = "group" + BLOCK = "block" diff --git a/python/tokenspeed/runtime/layers/quantization/base_config.py b/python/tokenspeed/runtime/layers/quantization/base_config.py new file mode 100755 index 0000000..063291d --- /dev/null +++ b/python/tokenspeed/runtime/layers/quantization/base_config.py @@ -0,0 +1,189 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + + +from abc import ABC, abstractmethod +from typing import Any + +import torch +from torch import nn + + +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 QuantizationConfig(ABC): + """Base class for quantization configs.""" + + def __init__( + self, + ignored_layers: list[str] | None = None, + exclude_modules: list[str] | None = None, + ) -> None: + self.ignored_layers = ignored_layers or [] + self.exclude_modules = exclude_modules or [] + + @abstractmethod + def get_name(self) -> str: + """Name of the quantization method.""" + raise NotImplementedError + + def moe_weight_dtype(self) -> str: + """Logical MoE weight dtype fed to ``moe_plan`` as the ``weight_dtype`` trait. + + Must name a concrete dtype the kernels register against (``fp8``, + ``nvfp4``, ``mxfp4``), not the quant method. Configs whose name already + is the dtype need no override; container formats (compressed-tensors) + resolve it from the parsed scheme. + """ + return self.get_name() + + @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., 90 for Hopper, 100 for Blackwell. + 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) -> str | None: + """ + 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_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 + + +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: torch.Tensor | None = 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 + + +def method_has_implemented_embedding(method_class: type[QuantizeMethodBase]) -> bool: + return "embedding" in method_class.__dict__ diff --git a/python/tokenspeed/runtime/layers/quantization/compressed_tensors/__init__.py b/python/tokenspeed/runtime/layers/quantization/compressed_tensors/__init__.py new file mode 100755 index 0000000..2b589d8 --- /dev/null +++ b/python/tokenspeed/runtime/layers/quantization/compressed_tensors/__init__.py @@ -0,0 +1,19 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. diff --git a/python/tokenspeed/runtime/layers/quantization/compressed_tensors/compressed_tensors.py b/python/tokenspeed/runtime/layers/quantization/compressed_tensors/compressed_tensors.py new file mode 100755 index 0000000..b221259 --- /dev/null +++ b/python/tokenspeed/runtime/layers/quantization/compressed_tensors/compressed_tensors.py @@ -0,0 +1,600 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import logging +from contextlib import suppress +from typing import Any, Literal, NamedTuple, cast + +import torch +from compressed_tensors.config import ( + CompressionFormat, + SparsityCompressionConfig, + SparsityStructure, +) +from compressed_tensors.quantization import ( + QuantizationArgs, + QuantizationStrategy, + QuantizationType, +) +from pydantic import BaseModel +from tokenspeed_kernel.platform import current_platform + +from tokenspeed.runtime.layers.quantization.base_config import ( + LinearMethodBase, + QuantizationConfig, +) +from tokenspeed.runtime.layers.quantization.compressed_tensors.gptq_marlin_moe import ( + is_activation_quantization_format, +) +from tokenspeed.runtime.layers.quantization.compressed_tensors.schemes import ( + WNA16_SUPPORTED_BITS, + CompressedTensorsScheme, + CompressedTensorsWNA16, +) +from tokenspeed.runtime.layers.quantization.utils import find_matched_target + +# ruff: noqa: F821 + + +logger = logging.getLogger(__name__) + +__all__ = ["CompressedTensorsLinearMethod"] + +SPARSITY_CONFIG_NAME: Literal["sparsity_config"] = "sparsity_config" +QUANTIZATION_SCHEME_MAP_TYPE = dict[str, dict[str, QuantizationArgs] | None] + + +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. + """ + if not 0 <= self.minor < 10: + raise ValueError(f"Invalid device capability minor version: {self.minor}.") + return self.major * 10 + self.minor + + +class CompressedTensorsConfig(QuantizationConfig): + DeepSeekFP8Config = None + + 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: dict[str, Any] | None = None, + config: dict[str, Any] | None = None, + packed_modules_mapping: dict[str, list[str]] | None = None, + ): + super().__init__(ignored_layers=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 + _packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} + self.packed_modules_mapping = packed_modules_mapping or _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 moe_weight_dtype(self) -> str: + # Container format: resolve the routed-expert scheme to a concrete MoE + # kernel dtype. Only INT4 group-32 symmetric pack-quantized weights + # (Kimi-K2.5 / K2.6 / K2.7, weight-only + bf16 group scales) are wired. + weight_quant = self.target_scheme_map["Linear"].get("weights") + input_quant = self.target_scheme_map["Linear"].get("input_activations") + if ( + weight_quant is not None + and self._is_wNa16_group_channel(weight_quant, input_quant) + and weight_quant.type == QuantizationType.INT + and weight_quant.num_bits == 4 + and weight_quant.strategy == QuantizationStrategy.GROUP.value + and weight_quant.group_size == 32 + and not weight_quant.actorder + ): + return "mxint4" + raise ValueError( + f"unsupported compressed-tensors MoE scheme for kernel selection: " + f"{weight_quant}" + ) + + def get_scaled_act_names(self) -> list[str]: + return [] + + @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: + if ( + target_scheme_map[target]["weights"].type + != QuantizationType.FLOAT + ): + raise ValueError( + "Activation quantization config is missing input_activations." + ) + 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: + platform = current_platform() + capability_tuple = DeviceCapability( + platform.arch_version.major, platform.arch_version.minor + ) + + 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_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 ( + 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, + ) + else: + raise ImportError( + "Other method (CompressedTensorsW4A16Sparse24) is not supported now" + ) + + 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): + 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, + ) + + raise NotImplementedError("No compressed-tensors compatible scheme was found.") + + def get_scheme( + self, layer: torch.nn.Module, layer_name: str | None = None + ) -> CompressedTensorsScheme | None: + """ + 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. + # 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: SparsityCompressionConfig | None = 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, + ): + raise ImportError("CompressedTensors24 is not supported now") + elif weight_quant is None: + logger.warning( + "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) -> str | None: + """ + 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 TokenSpeed + + :param name: param name + :return: matching param name for KV cache scale in TokenSpeed + """ + 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: QuantizationArgs | None, + input_quant: QuantizationArgs | None, + sparsity_scheme: SparsityCompressionConfig | None = 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, + ] + + if weight_quant is None or input_quant is None: + raise RuntimeError("Quantization args should be populated at this point.") + 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: torch.Tensor | None = 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/tokenspeed/runtime/layers/quantization/compressed_tensors/gptq_marlin_moe.py b/python/tokenspeed/runtime/layers/quantization/compressed_tensors/gptq_marlin_moe.py new file mode 100755 index 0000000..356e4ad --- /dev/null +++ b/python/tokenspeed/runtime/layers/quantization/compressed_tensors/gptq_marlin_moe.py @@ -0,0 +1,100 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import torch +from compressed_tensors import CompressionFormat + + +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 gptq_marlin_moe_repack( + b_q_weight: torch.Tensor, + perm: torch.Tensor, + size_k: int, + size_n: int, + num_bits: int, +) -> torch.Tensor: + from tokenspeed_kernel.ops.quantization.cuda import gptq_marlin_repack + + 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 + + +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_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 diff --git a/python/tokenspeed/runtime/layers/quantization/compressed_tensors/scalar_type.py b/python/tokenspeed/runtime/layers/quantization/compressed_tensors/scalar_type.py new file mode 100755 index 0000000..c90b9a7 --- /dev/null +++ b/python/tokenspeed/runtime/layers/quantization/compressed_tensors/scalar_type.py @@ -0,0 +1,371 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import functools +import struct +from dataclasses import dataclass +from enum import Enum + +_SCALAR_TYPES_ID_MAP = {} + + +class NanRepr(Enum): + NONE = 0 # nans are not supported + IEEE_754 = 1 # nans are: Exp all 1s, mantissa not all 0s + EXTD_RANGE_MAX_MIN = 2 # nans are: Exp all 1s, mantissa all 1s + + +# This ScalarType class is a parallel implementation of the C++ ScalarType +# class found in csrc/core/scalar_type.hpp. These two classes should be kept +# in sync until the inductor fully supports custom C++ classes. +@dataclass(frozen=True) +class ScalarType: + """ + ScalarType can represent a wide range of floating point and integer + types, in particular it can be used to represent sub-byte data types + (something that torch.dtype currently does not support). It is also + capable of representing types with a bias, i.e.: + `stored_value = value + bias`, + this is useful for quantized types (e.g. standard GPTQ 4bit uses a bias + of 8). The implementation for this class can be found in + csrc/core/scalar_type.hpp, these type signatures should be kept in sync + with that file. + """ + + exponent: int + """ + Number of bits in the exponent if this is a floating point type + (zero if this an integer type) + """ + + mantissa: int + """ + Number of bits in the mantissa if this is a floating point type, + or the number bits representing an integer excluding the sign bit if + this an integer type. + """ + + signed: bool + "If the type is signed (i.e. has a sign bit)" + + bias: int + """ + bias used to encode the values in this scalar type + (value = stored_value - bias, default 0) for example if we store the + type as an unsigned integer with a bias of 128 then the value 0 will be + stored as 128 and -1 will be stored as 127 and 1 will be stored as 129. + """ + + _finite_values_only: bool = False + """ + Private: if infs are supported, used `has_infs()` instead. + """ + + nan_repr: NanRepr = NanRepr.IEEE_754 + """ + How NaNs are represent in this scalar type, returns NanRepr value. + (not applicable for integer types) + """ + + def _floating_point_max_int(self) -> int: + assert ( + self.mantissa <= 52 and self.exponent <= 11 + ), f"Cannot represent max/min as a double for type {self.__str__()}" + + max_mantissa = (1 << self.mantissa) - 1 + if self.nan_repr == NanRepr.EXTD_RANGE_MAX_MIN: + max_mantissa = max_mantissa - 1 + + max_exponent = (1 << self.exponent) - 2 + if self.nan_repr == NanRepr.EXTD_RANGE_MAX_MIN or self.nan_repr == NanRepr.NONE: + assert ( + self.exponent < 11 + ), f"Cannot represent max/min as a double for type {self.__str__()}" + max_exponent = max_exponent + 1 + + # adjust the exponent to match that of a double + # for now we assume the exponent bias is the standard 2^(e-1) -1, (where + # e is the exponent bits), there is some precedent for non-standard + # biases, example `float8_e4m3b11fnuz` here: + # https://github.com/jax-ml/ml_dtypes but to avoid premature over + # complication we are just assuming the standard exponent bias until + # there is a need to support non-standard biases + exponent_bias = (1 << (self.exponent - 1)) - 1 + exponent_bias_double = (1 << 10) - 1 # double e = 11 + + max_exponent_double = max_exponent - exponent_bias + exponent_bias_double + + # shift the mantissa and exponent into the proper positions for an + # IEEE double and bitwise-or them together. + return (max_mantissa << (52 - self.mantissa)) | (max_exponent_double << 52) + + def _floating_point_max(self) -> float: + double_raw = self._floating_point_max_int() + return struct.unpack("!d", struct.pack("!Q", double_raw))[0] + + def _raw_max(self) -> int | float: + if self.is_floating_point(): + return self._floating_point_max() + else: + assert ( + self.size_bits < 64 or self.size_bits == 64 and self.is_signed() + ), "Cannot represent max as an int" + return (1 << self.mantissa) - 1 + + def _raw_min(self) -> int | float: + if self.is_floating_point(): + assert ( + self.is_signed() + ), "We currently assume all floating point types are signed" + sign_bit_double = 1 << 63 + + max_raw = self._floating_point_max_int() + min_raw = max_raw | sign_bit_double + return struct.unpack("!d", struct.pack("!Q", min_raw))[0] + else: + assert ( + not self.is_signed() or self.size_bits <= 64 + ), "Cannot represent min as a int64_t" + + if self.is_signed(): + return -(1 << (self.size_bits - 1)) + else: + return 0 + + @functools.cached_property + def id(self) -> int: + """ + Convert the ScalarType to an int which can be passed to pytorch custom + ops. This layout of the int must be kept in sync with the C++ + ScalarType's from_id method. + """ + val = 0 + offset = 0 + + def or_and_advance(member, bit_width): + nonlocal val + nonlocal offset + bit_mask = (1 << bit_width) - 1 + val = val | (int(member) & bit_mask) << offset + offset = offset + bit_width + + or_and_advance(self.exponent, 8) + or_and_advance(self.mantissa, 8) + or_and_advance(self.signed, 1) + or_and_advance(self.bias, 32) + or_and_advance(self._finite_values_only, 1) + or_and_advance(self.nan_repr.value, 8) + + assert offset <= 64, f"ScalarType fields too big {offset} to fit into an int64" + + _SCALAR_TYPES_ID_MAP[val] = self + + return val + + @property + def size_bits(self) -> int: + return self.exponent + self.mantissa + int(self.signed) + + def min(self) -> int | float: + """ + Min representable value for this scalar type. + (accounting for bias if there is one) + """ + return self._raw_min() - self.bias + + def max(self) -> int | float: + """ + Max representable value for this scalar type. + (accounting for bias if there is one) + """ + return self._raw_max() - self.bias + + def is_signed(self) -> bool: + """ + If the type is signed (i.e. has a sign bit), same as `signed` + added for consistency with: + https://pytorch.org/docs/stable/generated/torch.Tensor.is_signed.html + """ + return self.signed + + def is_floating_point(self) -> bool: + "If the type is a floating point type" + return self.exponent != 0 + + def is_integer(self) -> bool: + "If the type is an integer type" + return self.exponent == 0 + + def has_bias(self) -> bool: + "If the type has a non-zero bias" + return self.bias != 0 + + def has_infs(self) -> bool: + "If the type is floating point and supports infinity" + return not self._finite_values_only + + def has_nans(self) -> bool: + return self.nan_repr != NanRepr.NONE + + def is_ieee_754(self) -> bool: + """ + If the type is a floating point type that follows IEEE 754 + conventions + """ + return self.nan_repr == NanRepr.IEEE_754 and not self._finite_values_only + + def __str__(self) -> str: + """ + naming generally follows: https://github.com/jax-ml/ml_dtypes + for floating point types (leading f) the scheme is: + `float_em[flags]` + flags: + - no-flags: means it follows IEEE 754 conventions + - f: means finite values only (no infinities) + - n: means nans are supported (non-standard encoding) + for integer types the scheme is: + `[u]int[b]` + - if bias is not present it means its zero + """ + if self.is_floating_point(): + ret = ( + "float" + + str(self.size_bits) + + "_e" + + str(self.exponent) + + "m" + + str(self.mantissa) + ) + + if not self.is_ieee_754(): + if self._finite_values_only: + ret = ret + "f" + if self.nan_repr != NanRepr.NONE: + ret = ret + "n" + + return ret + else: + ret = ("int" if self.is_signed() else "uint") + str(self.size_bits) + if self.has_bias(): + ret = ret + "b" + str(self.bias) + return ret + + def __repr__(self) -> str: + return "ScalarType." + self.__str__() + + # __len__ needs to be defined (and has to throw TypeError) for pytorch's + # opcheck to work. + def __len__(self) -> int: + raise TypeError + + # + # Convenience Constructors + # + + @classmethod + def int_(cls, size_bits: int, bias: int | None) -> "ScalarType": + "Create a signed integer scalar type (size_bits includes sign-bit)." + ret = cls(0, size_bits - 1, True, bias if bias else 0) + ret.id # noqa B018: make sure the id is cached + return ret + + @classmethod + def uint(cls, size_bits: int, bias: int | None) -> "ScalarType": + """Create a unsigned integer scalar type.""" + ret = cls(0, size_bits, False, bias if bias else 0) + ret.id # noqa B018: make sure the id is cached + return ret + + @classmethod + def float_IEEE754(cls, exponent: int, mantissa: int) -> "ScalarType": + """ + Create a standard floating point type + (i.e. follows IEEE 754 conventions). + """ + assert mantissa > 0 and exponent > 0 + ret = cls(exponent, mantissa, True, 0) + ret.id # noqa B018: make sure the id is cached + return ret + + @classmethod + def float_( + cls, exponent: int, mantissa: int, finite_values_only: bool, nan_repr: NanRepr + ) -> "ScalarType": + """ + Create a non-standard floating point type + (i.e. does not follow IEEE 754 conventions). + """ + assert mantissa > 0 and exponent > 0 + assert nan_repr != NanRepr.IEEE_754, ( + "use `float_IEEE754` constructor for floating point types that " + "follow IEEE 754 conventions" + ) + ret = cls(exponent, mantissa, True, 0, finite_values_only, nan_repr) + ret.id # noqa B018: make sure the id is cached + return ret + + @classmethod + def from_id(cls, scalar_type_id: int): + if scalar_type_id not in _SCALAR_TYPES_ID_MAP: + raise ValueError(f"scalar_type_id {scalar_type_id} doesn't exists.") + return _SCALAR_TYPES_ID_MAP[scalar_type_id] + + +# naming generally follows: https://github.com/jax-ml/ml_dtypes +# for floating point types (leading f) the scheme is: +# `float_em[flags]` +# flags: +# - no-flags: means it follows IEEE 754 conventions +# - f: means finite values only (no infinities) +# - n: means nans are supported (non-standard encoding) +# for integer types the scheme is: +# `[u]int[b]` +# - if bias is not present it means its zero + + +class scalar_types: + int4 = ScalarType.int_(4, None) + uint4 = ScalarType.uint(4, None) + int8 = ScalarType.int_(8, None) + uint8 = ScalarType.uint(8, None) + float8_e4m3fn = ScalarType.float_(4, 3, True, NanRepr.EXTD_RANGE_MAX_MIN) + float8_e5m2 = ScalarType.float_IEEE754(5, 2) + float16_e8m7 = ScalarType.float_IEEE754(8, 7) + float16_e5m10 = ScalarType.float_IEEE754(5, 10) + + # fp6, https://github.com/usyd-fsalab/fp6_llm/tree/main + float6_e3m2f = ScalarType.float_(3, 2, True, NanRepr.NONE) + + # fp4, https://www.opencompute.org/documents/ocp-microscaling-formats-mx-v1-0-spec-final-pdf + float4_e2m1f = ScalarType.float_(2, 1, True, NanRepr.NONE) + + # "gptq" types + uint2b2 = ScalarType.uint(2, 2) + uint3b4 = ScalarType.uint(3, 4) + uint4b8 = ScalarType.uint(4, 8) + uint8b128 = ScalarType.uint(8, 128) + + # colloquial names + bfloat16 = float16_e8m7 + float16 = float16_e5m10 diff --git a/python/tokenspeed/runtime/layers/quantization/compressed_tensors/schemes/__init__.py b/python/tokenspeed/runtime/layers/quantization/compressed_tensors/schemes/__init__.py new file mode 100755 index 0000000..ee8369d --- /dev/null +++ b/python/tokenspeed/runtime/layers/quantization/compressed_tensors/schemes/__init__.py @@ -0,0 +1,33 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from tokenspeed.runtime.layers.quantization.compressed_tensors.schemes.compressed_tensors_scheme import ( + CompressedTensorsScheme, +) +from tokenspeed.runtime.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( + WNA16_SUPPORTED_BITS, + CompressedTensorsWNA16, +) + +__all__ = [ + "CompressedTensorsScheme", + "CompressedTensorsWNA16", + "WNA16_SUPPORTED_BITS", +] diff --git a/python/tokenspeed/runtime/layers/quantization/compressed_tensors/schemes/compressed_tensors_scheme.py b/python/tokenspeed/runtime/layers/quantization/compressed_tensors/schemes/compressed_tensors_scheme.py new file mode 100755 index 0000000..9ceb191 --- /dev/null +++ b/python/tokenspeed/runtime/layers/quantization/compressed_tensors/schemes/compressed_tensors_scheme.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + + +from abc import ABC, abstractmethod + +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: torch.Tensor | None + ): + """ + 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/tokenspeed/runtime/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py b/python/tokenspeed/runtime/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py new file mode 100755 index 0000000..f00203b --- /dev/null +++ b/python/tokenspeed/runtime/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16.py @@ -0,0 +1,339 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +# ruff: noqa: F821 + +import logging +from collections.abc import Callable + +import torch +from compressed_tensors.quantization import ActivationOrdering +from tokenspeed_kernel.ops.quantization.cuda import gptq_marlin_repack + +# yapf conflicts with isort for this block +# yapf: disable +from tokenspeed.runtime.layers.parameter import ( + BaseWeightParameter, + ChannelQuantScaleParameter, + GroupQuantScaleParameter, + PackedColumnParameter, + PackedWeightParameter, + RowParallelWeightParameter, + permute_param_layout_, +) +from tokenspeed.runtime.layers.quantization.compressed_tensors.scalar_type import ( + scalar_types, +) +from tokenspeed.runtime.layers.quantization.compressed_tensors.schemes import ( + CompressedTensorsScheme, +) +from tokenspeed.runtime.layers.quantization.utils import replace_parameter, unpack_cols + +logger = logging.getLogger(__name__) + +__all__ = ["CompressedTensorsWNA16"] +WNA16_SUPPORTED_TYPES_MAP = { + 4: scalar_types.uint4b8, + 8: scalar_types.uint8b128 +} +WNA16_ZP_SUPPORTED_TYPES_MAP = {4: scalar_types.uint4, 8: scalar_types.uint8} +WNA16_SUPPORTED_BITS = list(WNA16_SUPPORTED_TYPES_MAP.keys()) + + +class CompressedTensorsWNA16(CompressedTensorsScheme): + _kernel_backends_being_used: set[str] = set() + + def __init__(self, + strategy: str, + num_bits: int, + group_size: int | None = None, + symmetric: bool | None = True, + actorder: ActivationOrdering | None = None): + + self.pack_factor = 32 // num_bits + self.strategy = strategy + self.symmetric = symmetric + self.group_size = -1 if group_size is None else group_size + self.has_g_idx = actorder == ActivationOrdering.GROUP + + if self.group_size == -1 and self.strategy != "channel": + raise ValueError("Marlin kernels require group quantization or " + "channelwise quantization, but found no group " + "size and strategy is not channelwise.") + + if num_bits not in WNA16_SUPPORTED_TYPES_MAP: + raise ValueError( + f"Unsupported num_bits = {num_bits}. " + f"Supported num_bits = {WNA16_SUPPORTED_TYPES_MAP.keys()}") + + self.quant_type = (WNA16_ZP_SUPPORTED_TYPES_MAP[num_bits] + if not self.symmetric else + WNA16_SUPPORTED_TYPES_MAP[num_bits]) + + @classmethod + def get_min_capability(cls) -> int: + return 90 + + def create_weights(self, layer: torch.nn.Module, output_size: int, + 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) + + 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_type, + act_type=params_dtype, + group_size=self.group_size, + zero_points=not self.symmetric, + has_g_idx=self.has_g_idx + ) + + # If group_size is -1, we are in channelwise case. + group_size = self.group_size if self.group_size != -1 else input_size + row_parallel = (input_size != input_size_per_partition) + partition_scales = not marlin_repeat_scales_on_all_ranks( + self.has_g_idx, self.group_size, row_parallel) + + scales_and_zp_size = input_size // group_size + + if partition_scales: + assert input_size_per_partition % group_size == 0 + scales_and_zp_size = input_size_per_partition // group_size + + weight = PackedWeightParameter(input_dim=1, + output_dim=0, + weight_loader=weight_loader, + packed_factor=self.pack_factor, + packed_dim=1, + data=torch.empty( + output_size_per_partition, + input_size_per_partition // + self.pack_factor, + dtype=torch.int32, + )) + + weight_scale_args = { + "weight_loader": + weight_loader, + "data": + torch.empty( + output_size_per_partition, + scales_and_zp_size, + dtype=params_dtype, + ) + } + + zeros_args = { + "weight_loader": + weight_loader, + "data": + torch.zeros( + output_size_per_partition // self.pack_factor, + scales_and_zp_size, + dtype=torch.int32, + ) + } + + if not partition_scales: + weight_scale = ChannelQuantScaleParameter(output_dim=0, + **weight_scale_args) + + if not self.symmetric: + qzeros = PackedColumnParameter(output_dim=0, + packed_dim=0, + packed_factor=self.pack_factor, + **zeros_args) + else: + weight_scale = GroupQuantScaleParameter(output_dim=0, + input_dim=1, + **weight_scale_args) + if not self.symmetric: + qzeros = PackedWeightParameter(input_dim=1, + output_dim=0, + packed_dim=0, + packed_factor=self.pack_factor, + **zeros_args) + + # A 2D array defining the original shape of the weights + # before packing + weight_shape = BaseWeightParameter(data=torch.empty(2, + dtype=torch.int64), + weight_loader=weight_loader) + + layer.register_parameter("weight_packed", weight) + layer.register_parameter("weight_scale", weight_scale) + layer.register_parameter("weight_shape", weight_shape) + + if not self.symmetric: + layer.register_parameter("weight_zero_point", qzeros) + + # group index (for activation reordering) + if self.has_g_idx: + weight_g_idx = RowParallelWeightParameter(data=torch.empty( + input_size_per_partition, + dtype=torch.int32, + ), + input_dim=0, + weight_loader=weight_loader) + layer.register_parameter("weight_g_idx", weight_g_idx) + + # Checkpoints are serialized in compressed-tensors format, which is + # different from the format the kernel may want. Handle repacking here. + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + # Default names since marlin requires empty parameters for these, + self.w_q_name = "weight_packed" + self.w_s_name = "weight_scale" + self.w_zp_name = "weight_zero_point" + self.w_gidx_name = "weight_g_idx" + + device = getattr(layer, self.w_q_name).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) + + def _transform_param( + layer: torch.nn.Module, name: str | None, 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, BaseWeightParameter) + 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, BaseWeightParameter) + 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_weights(self, layer: torch.nn.Module, x: torch.Tensor, + bias: 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 + torch.Tensor | None, # w_zp, + torch.Tensor | None, # 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, + ) diff --git a/python/tokenspeed/runtime/layers/quantization/fp8.py b/python/tokenspeed/runtime/layers/quantization/fp8.py new file mode 100755 index 0000000..24e054e --- /dev/null +++ b/python/tokenspeed/runtime/layers/quantization/fp8.py @@ -0,0 +1,104 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + + +from __future__ import annotations + +import logging +from typing import Any + +import torch + +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.utils import log_info_on_rank0 + +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: list[str] | None = None, + weight_block_size: list[int] = None, + scale_fmt: str | None = None, + ) -> None: + super().__init__(ignored_layers=ignored_layers) + 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 + if weight_block_size is not None: + if not is_checkpoint_fp8_serialized: + raise ValueError( + "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 + self.scale_fmt = scale_fmt.lower() if scale_fmt is not None else None + + @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 90 + + @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) + scale_fmt = cls.get_from_keys_or(config, ["scale_fmt"], 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, + scale_fmt=scale_fmt, + ) + + def get_scaled_act_names(self) -> list[str]: + return [] diff --git a/python/tokenspeed/runtime/layers/quantization/mxfp4.py b/python/tokenspeed/runtime/layers/quantization/mxfp4.py new file mode 100644 index 0000000..d2fd355 --- /dev/null +++ b/python/tokenspeed/runtime/layers/quantization/mxfp4.py @@ -0,0 +1,205 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + + +from __future__ import annotations + +import re +from collections.abc import Mapping +from typing import Any + +import torch +from tokenspeed_kernel.platform import current_platform + +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig + + +def _is_fp4_e8m0_per_group(stage: object, *, is_dynamic: bool | None = None) -> bool: + if not isinstance(stage, Mapping): + return False + if is_dynamic is not None and stage.get("is_dynamic") is not is_dynamic: + return False + return ( + str(stage.get("dtype", "")).lower() in {"fp4", "mxfp4"} + and str(stage.get("qscheme", "")).lower() == "per_group" + and stage.get("group_size") in {32, "32"} + and str(stage.get("scale_format", "")).lower() == "e8m0" + ) + + +def _is_amd_quark_w_mxfp4_a_fp8(config: Mapping[str, Any]) -> bool: + if not isinstance(config, Mapping): + return False + if not current_platform().is_amd: + return False + if str(config.get("quant_method", "")).lower() != "quark": + return False + global_quant_config = config.get("global_quant_config") or {} + export = config.get("export") or {} + if not isinstance(global_quant_config, Mapping) or not isinstance(export, Mapping): + return False + input_tensors = global_quant_config.get("input_tensors") or {} + weight = global_quant_config.get("weight") or {} + return ( + isinstance(input_tensors, Mapping) + and "fp8" in str(input_tensors.get("dtype", "")).lower() + and _is_fp4_e8m0_per_group(weight, is_dynamic=False) + and str(export.get("pack_method", "")).lower() == "reorder" + and str(export.get("weight_format", "")).lower() == "real_quantized" + ) + + +def _is_amd_quark_dynamic_mxfp4(config: Mapping[str, Any]) -> bool: + if not isinstance(config, Mapping): + return False + if not current_platform().is_amd: + return False + if str(config.get("quant_method", "")).lower() != "quark": + return False + global_quant_config = config.get("global_quant_config") or {} + export = config.get("export") or {} + if not isinstance(global_quant_config, Mapping) or not isinstance(export, Mapping): + return False + input_tensors = global_quant_config.get("input_tensors") or {} + weight = global_quant_config.get("weight") or {} + return ( + _is_fp4_e8m0_per_group(input_tensors, is_dynamic=True) + and _is_fp4_e8m0_per_group(weight, is_dynamic=False) + and str(export.get("pack_method", "")).lower() == "reorder" + and str(export.get("weight_format", "")).lower() == "real_quantized" + ) + + +def _is_amd_quark_mxfp4_checkpoint(config: dict) -> bool: + if not isinstance(config, Mapping): + return False + return _is_amd_quark_w_mxfp4_a_fp8(config) or _is_amd_quark_dynamic_mxfp4(config) + + +def _iter_ignored_layer_pattern_aliases(raw: str): + yield raw + if raw.startswith("language_model."): + yield raw.removeprefix("language_model.") + return + + if raw.startswith("re:"): + regex = raw[3:] + for prefix in ("language_model.", re.escape("language_model.")): + if regex.startswith(prefix): + yield f"re:{regex.removeprefix(prefix)}" + return + + +def _to_ignore_pattern(raw: str) -> str: + if raw.startswith("re:") or "*" not in raw: + return raw + regex = re.escape(raw).replace(r"\*", ".*") + return f"re:{regex}" + + +def _normalize_ignored_layer_patterns(patterns: list[str] | None) -> list[str]: + """Normalize ignored-layer patterns into the form understood by + ``should_ignore_quant_layer``. + + Some exporters (notably AMD-Quark) accept shell-style globs such as + ``"*lm_head"`` or ``"*self_attn*"``. ``should_ignore_quant_layer`` + expects either an exact name or a regex prefixed with ``re:``. Convert + glob-like entries to regex while passing through plain literals. + """ + if not patterns: + return [] + normalized: list[str] = [] + seen: set[str] = set() + for raw in patterns: + if not isinstance(raw, str) or not raw: + continue + for alias in _iter_ignored_layer_pattern_aliases(raw): + pattern = _to_ignore_pattern(alias) + if pattern in seen: + continue + seen.add(pattern) + normalized.append(pattern) + return normalized + + +class Mxfp4Config(QuantizationConfig): + + def __init__( + self, + ignored_layers: list[str] | None = None, + is_checkpoint_mxfp4_serialized: bool = False, + is_w4a8_fp8: bool = False, + use_dynamic_mxfp4_activations: bool = False, + ): + super().__init__(ignored_layers=ignored_layers) + self.is_checkpoint_mxfp4_serialized = is_checkpoint_mxfp4_serialized + self.is_w4a8_fp8 = is_w4a8_fp8 + self.use_dynamic_mxfp4_activations = use_dynamic_mxfp4_activations + self.group_size = 32 + + @classmethod + def from_config(cls, config): + quant_method = str(config.get("quant_method", "")).lower() + is_w4a8_fp8 = _is_amd_quark_w_mxfp4_a_fp8(config) + use_dynamic_mxfp4_activations = _is_amd_quark_dynamic_mxfp4(config) + is_checkpoint_mxfp4_serialized = ( + "mxfp4" in quant_method or is_w4a8_fp8 or use_dynamic_mxfp4_activations + ) + + raw_ignored = cls.get_from_keys_or(config, ["ignored_layers", "exclude"], None) + ignored_layers = _normalize_ignored_layer_patterns(raw_ignored) + + return cls( + ignored_layers=ignored_layers, + is_checkpoint_mxfp4_serialized=is_checkpoint_mxfp4_serialized, + is_w4a8_fp8=is_w4a8_fp8, + use_dynamic_mxfp4_activations=use_dynamic_mxfp4_activations, + ) + + @classmethod + def override_quantization_method(cls, hf_quant_cfg, user_quant) -> str | None: + """Promote AMD Quark MXFP4 checkpoint metadata to mxfp4.""" + if user_quant in {"mxfp4", None} and _is_amd_quark_mxfp4_checkpoint( + hf_quant_cfg + ): + return "mxfp4" + return None + + @classmethod + def get_min_capability(cls) -> int: + return 90 + + @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_scaled_act_names(self) -> list[str]: + return [] diff --git a/python/tokenspeed/runtime/layers/quantization/nvfp4.py b/python/tokenspeed/runtime/layers/quantization/nvfp4.py new file mode 100644 index 0000000..2a1c297 --- /dev/null +++ b/python/tokenspeed/runtime/layers/quantization/nvfp4.py @@ -0,0 +1,115 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""NVFP4 quantization config for tokenspeed runtime (ModelOpt-produced checkpoints).""" + +import logging +from typing import Any + +import torch + +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig + +logger = logging.getLogger(__name__) + + +class Nvfp4Config(QuantizationConfig): + """Config class for NVFP4 quantization (ModelOpt-produced checkpoints).""" + + def __init__( + self, + kv_cache_quant_algo: str | None = None, + group_size: int = 16, + exclude_modules: list[str] | None = None, + ) -> None: + super().__init__(exclude_modules=exclude_modules) + self.kv_cache_quant_algo = kv_cache_quant_algo + self.group_size = group_size + self.weight_block_size = None # FP4 uses group_size, not weight_block_size + + @classmethod + def get_name(cls) -> str: + return "nvfp4" + + @classmethod + def get_supported_act_dtypes(cls) -> list[torch.dtype]: + return [torch.bfloat16, torch.half] + + @classmethod + def get_min_capability(cls) -> int: + return 100 # Blackwell required + + @staticmethod + def get_config_filenames() -> list[str]: + return ["hf_quant_config.json"] + + @classmethod + def from_config(cls, config: dict[str, Any]) -> "Nvfp4Config": + kv_cache_quant_algo = None + group_size = 16 + exclude_modules = [] + + # Try flat format first (config.json quantization_config) + quant_method = config.get("quant_algo") + if quant_method is not None: + kv_cache_quant_algo = config.get("kv_cache_quant_algo", "auto") + group_size = config.get("group_size", 16) + exclude_modules = config.get("ignore", []) + else: + # Fall back to nested format (hf_quant_config.json) + 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", "auto") + group_size = quant_config.get("group_size", 16) + exclude_modules = quant_config.get("exclude_modules", []) + except (ValueError, KeyError): + raise ValueError( + "Cannot find quant_algo in the model quantization config." + ) + + if quant_method != "NVFP4": + raise ValueError(f"Nvfp4Config only supports NVFP4, got {quant_method}") + + return cls( + kv_cache_quant_algo=kv_cache_quant_algo, + group_size=group_size, + exclude_modules=exclude_modules, + ) + + @classmethod + def override_quantization_method(cls, hf_quant_cfg, user_quant) -> str | None: + """Detect NVFP4 from hf_quant_config and override.""" + quant_algo = "" + if isinstance(hf_quant_cfg, dict): + quant_algo = hf_quant_cfg.get("quant_algo", "") + if not quant_algo: + q = hf_quant_cfg.get("quantization", {}) + if isinstance(q, dict): + quant_algo = q.get("quant_algo", "") + if "NVFP4" in quant_algo.upper() or "FP4" in quant_algo.upper(): + return "nvfp4" + # Fallback: user requested nvfp4 and the checkpoint was produced by ModelOpt. + if user_quant == "nvfp4" and hf_quant_cfg.get("quant_method") == "modelopt": + return "nvfp4" + return None + + def get_scaled_act_names(self) -> list[str]: + return [] diff --git a/python/tokenspeed/runtime/layers/quantization/utils.py b/python/tokenspeed/runtime/layers/quantization/utils.py new file mode 100755 index 0000000..b6aac5f --- /dev/null +++ b/python/tokenspeed/runtime/layers/quantization/utils.py @@ -0,0 +1,404 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import re +from collections.abc import Iterable, Mapping +from types import MappingProxyType + +import numpy +import torch +from torch.nn import Module + +from tokenspeed.runtime.layers.quantization.compressed_tensors.scalar_type import ( + ScalarType as ScalarType, +) + + +def should_exclude_quant_module(prefix: str, exclude_modules: list[str]) -> bool: + """Whether ``prefix`` matches a ModelOpt-style glob in ``exclude_modules``.""" + if prefix is None or not exclude_modules: + return False + for pattern in exclude_modules: + regex_str = pattern.replace(".", r"\.").replace("*", ".*") + if re.fullmatch(regex_str, prefix): + return True + return False + + +def should_ignore_quant_layer( + prefix: str, + ignored_layers: list[str], + fused_mapping: Mapping[str, list[str]] = MappingProxyType({}), +) -> bool: + if prefix is None or ignored_layers is None: + return False + + # layer_name = model.layers.0.self_attn.qkv_proj + # proj_name = qkv_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 and prefix not in ignored_layers: + shard_proj_names = fused_mapping[proj_name] + + # Convert fused_name --> [shard_names] + shard_names = [ + prefix.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=ignored_layers + ) + + # 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 {prefix}. TokenSpeed " + "requires all to use the same scheme." + ) + else: + should_ignore_layer = check_equal_or_regex_match( + layer_name=prefix, targets=ignored_layers + ) + if not should_ignore_layer: + 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: + should_ignore_layer = True + elif "fused_qkv_a_proj_with_mqa" in prefix: + prefix_q_a_proj = prefix.replace( + "fused_qkv_a_proj_with_mqa", "q_a_proj" + ) + prefix_kv_a_proj_with_mqa = prefix.replace( + "fused_qkv_a_proj_with_mqa", "kv_a_proj_with_mqa" + ) + if ( + prefix_q_a_proj in ignored_layers + and prefix_kv_a_proj_with_mqa in ignored_layers + ): + should_ignore_layer = True + elif "qkv_proj" in prefix: + prefix_q_proj = prefix.replace("qkv_proj", "q_proj") + prefix_k_proj = prefix.replace("qkv_proj", "k_proj") + prefix_v_proj = prefix.replace("qkv_proj", "v_proj") + if ( + prefix_q_proj in ignored_layers + and prefix_k_proj in ignored_layers + and prefix_v_proj in ignored_layers + ): + should_ignore_layer = True + elif "experts" in prefix: + should_ignore_layer = any( + [ + prefix in layer_name + for layer_name in ignored_layers + if "experts" in layer_name + ] + ) + + if should_ignore_layer is None: + raise RuntimeError("Layer ignore decision was not initialized.") + 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: str | None, + 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 +) -> str | None: + """ + 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]], +) -> str | None: + """ + 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[str | None] = [] + 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 + + +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 update_tensor_inplace(old: torch.Tensor, new: torch.Tensor) -> None: + old.copy_(new) + + +# Newly generated tensors need to replace existing tensors that are +# already registered as parameters by TokenSpeed (and won't be freed) +def replace_parameter( + mod: torch.nn.Module, name: str, new: 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 get_pack_factor(num_bits): + if num_bits <= 0 or 32 % num_bits != 0: + raise ValueError(f"Unsupported num_bits = {num_bits}") + return 32 // num_bits + + +def unpack_cols( + packed_q_w: torch.Tensor, + num_bits: int, + size_k: int, + size_n: int, +): + pack_factor = get_pack_factor(num_bits) + if size_n % pack_factor != 0: + raise ValueError(f"size_n={size_n} must be divisible by {pack_factor}.") + expected_shape = (size_k, size_n // pack_factor) + if packed_q_w.shape != expected_shape: + raise ValueError( + f"packed_q_w.shape = {packed_q_w.shape} size_k = {size_k}, " + f"size_n = {size_n} pack_Factor = {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 + + +def block_dequant( + x_q_block: torch.Tensor, + x_s: torch.Tensor, + block_size: list[int], +) -> tuple[torch.Tensor, torch.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 + if n_tiles != x_s.shape[0] or k_tiles != x_s.shape[1]: + raise ValueError( + f"Scale shape {tuple(x_s.shape)} does not match tiles " + f"({n_tiles}, {k_tiles})." + ) + + 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] + + return x_dq_block diff --git a/python/tokenspeed/runtime/layers/quantization/w8a8_fp8.py b/python/tokenspeed/runtime/layers/quantization/w8a8_fp8.py new file mode 100755 index 0000000..91d1422 --- /dev/null +++ b/python/tokenspeed/runtime/layers/quantization/w8a8_fp8.py @@ -0,0 +1,76 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from typing import Any + +import torch + +from tokenspeed.runtime.layers.quantization import QuantizationConfig + + +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): + super().__init__() + 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]): + 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_scaled_act_names(self) -> list[str]: + return [] diff --git a/python/tokenspeed/runtime/layers/rotary_embedding.py b/python/tokenspeed/runtime/layers/rotary_embedding.py new file mode 100755 index 0000000..47b12a0 --- /dev/null +++ b/python/tokenspeed/runtime/layers/rotary_embedding.py @@ -0,0 +1,1468 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + + +"""Rotary Positional Embeddings.""" + +import itertools +import logging +import math +from typing import Any + +import torch +import torch.nn as nn +from tokenspeed_kernel.ops.embedding import FusedSetKVBufferArg, apply_rope +from tokenspeed_kernel.platform import current_platform +from tokenspeed_kernel.torch_compile import get_compiler_backend + +_is_nvidia = current_platform().is_nvidia +logger = logging.getLogger(__name__) + + +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) + + +# 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) + + +@torch.compile(dynamic=True, backend=get_compiler_backend()) +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_interleaved_rope(x: torch.Tensor, mrope_section: list) -> torch.Tensor: + x_t = x[0].clone() + x_t[..., 1 : mrope_section[1] * 3 : 3] = x[1, ..., 1 : mrope_section[1] * 3 : 3] + x_t[..., 2 : mrope_section[2] * 3 : 3] = x[2, ..., 2 : mrope_section[2] * 3 : 3] + return x_t + + +class RotaryEmbedding(torch.nn.Module): + """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() + self.cos_sin_cache: torch.Tensor + self.register_buffer("cos_sin_cache", cache, persistent=False) + + def _compute_inv_freq(self, base: int | float) -> torch.Tensor: + """Compute the inverse frequency.""" + # 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( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + offsets: torch.Tensor | None = None, + fused_set_kv_buffer_arg: FusedSetKVBufferArg | None = None, + output_q_rope: torch.Tensor | None = None, + output_k_rope: torch.Tensor | None = None, + enable_pdl: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + if offsets is not None: + raise ValueError("embedding.rope does not support offsets") + return apply_rope( + positions=positions, + q=query, + k=key, + head_size=self.head_size, + cos_sin_cache=self.cos_sin_cache, + is_neox=self.is_neox_style, + fused_set_kv_buffer_arg=fused_set_kv_buffer_arg, + q_rope_out=output_q_rope, + k_rope_out=output_k_rope, + enable_pdl=enable_pdl, + ) + + 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: list[float] | float, + dtype: torch.dtype, + ) -> None: + if isinstance(scaling_factors, float): + scaling_factors = [scaling_factors] + self.scaling_factors: list[float] = scaling_factors + 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: + # 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) + } + if len(self.scaling_factors) != len(offsets): + raise RuntimeError("scaling factor offsets were not initialized correctly.") + 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: + # 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.""" + + 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: float | None = None, + long_mscale: float | None = 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: torch.Tensor | None = 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: str | None = "cuda", + ) -> 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 + ) + + 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( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + fused_set_kv_buffer_arg=None, + output_q_rope=None, + offsets: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if _is_nvidia: + return super().forward( + positions=positions, + query=query, + key=key, + fused_set_kv_buffer_arg=fused_set_kv_buffer_arg, + output_q_rope=output_q_rope, + offsets=offsets, + ) + + 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: + # 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) + + +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: 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 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: list[int] | None = None, + mrope_interleaved: bool = False, + ) -> None: + super().__init__( + head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype + ) + + self.mrope_section = mrope_section + self.mrope_interleaved = mrope_interleaved + if self.mrope_section: + expected_sum = rotary_dim // 2 + actual_sum = sum(self.mrope_section) + if actual_sum != expected_sum: + logger.warning( + 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 + + logger.warning( + f"Corrected mrope_section: {self.mrope_section} (sum={sum(self.mrope_section)})" + ) + + @torch.compile(dynamic=True, backend=get_compiler_backend()) + def forward( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor, + fused_set_kv_buffer_arg: FusedSetKVBufferArg | None = None, + ) -> 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] + """ + if fused_set_kv_buffer_arg is not None: + raise ValueError("save kv cache is not supported for MRotaryEmbedding.") + if positions.ndim not in (1, 2): + raise ValueError(f"positions must be 1D or 2D, got ndim={positions.ndim}.") + + num_tokens = positions.shape[-1] + cos_sin = self.cos_sin_cache[positions] + cos, sin = cos_sin.chunk(2, dim=-1) + if positions.ndim == 2: + if not self.mrope_section: + raise RuntimeError("mrope_section must be set for 2D M-RoPE.") + + if self.mrope_interleaved: + cos = apply_interleaved_rope(cos, self.mrope_section) + sin = apply_interleaved_rope(sin, self.mrope_section) + else: + 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 + + @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: int | None = None, + input_ids: torch.LongTensor | None = None, + image_grid_thw: torch.LongTensor | None = None, + video_grid_thw: torch.LongTensor | None = None, + second_per_grid_ts: torch.Tensor | None = 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 in ( + "qwen2_vl", + "qwen3_vl", + "qwen3_vl_moe", + "qwen3_5", + "qwen3_5_moe", + ): + 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 + + @staticmethod + def get_rope_index_glm4v( + input_ids: torch.Tensor, + hf_config: Any, + image_grid_thw: list[list[int]] | torch.Tensor, + video_grid_thw: 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 + + +_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: dict[str, Any] | None = None, + dtype: torch.dtype | None = None, + partial_rotary_factor: float = 1.0, +) -> 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] + + if 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"], + mrope_interleaved=rope_scaling.get("mrope_interleaved", False), + ) + 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"] + 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 diff --git a/python/tokenspeed/runtime/layers/utils.py b/python/tokenspeed/runtime/layers/utils.py new file mode 100755 index 0000000..addc13e --- /dev/null +++ b/python/tokenspeed/runtime/layers/utils.py @@ -0,0 +1,111 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Utilities for context-parallel layer helpers.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass + +import torch + +from tokenspeed.runtime.distributed.comm_ops import token_all_gather +from tokenspeed.runtime.utils import get_bool_env_var + + +def get_layer_id(weight_name: str) -> int | None: + # 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 + + +# Attention CP utils +@dataclass +class ContextParallelMetadata: + split_list: list[int] | None = None + inverse_split_list: list[int] | None = None + max_token_len_in_block: int = -1 + zigzag_index: list[int] | None = None + per_rank_actual_token: list[int] | None = None + prefix_sum_tokens_prev: int = -1 + prefix_sum_tokens_cur: int = -1 + tokens_prev: int = -1 + tokens_cur: int = -1 + total_token_len: int = -1 + + +class CPMetadataContainer: + """Container for storing global CP metadata.""" + + def __init__(self): + self.value: ContextParallelMetadata | None = None + + def set(self, metadata: ContextParallelMetadata | None) -> None: + self.value = metadata + + def get(self) -> ContextParallelMetadata | None: + return self.value + + def __bool__(self) -> bool: + """Support ``if CP_METADATA`` syntax.""" + return self.value is not None + + +CP_METADATA = CPMetadataContainer() +ENABLE_CP = get_bool_env_var("ENABLE_CP", "false") + + +def cp_split_and_rebuild_data(x: torch.Tensor, split_list, zigzag_index): + split_tensors = list(torch.split(x, split_list, dim=0)) + return torch.cat([split_tensors[i] for i in zigzag_index], dim=0) + + +def cp_all_gather_rerange_output( + x, cp_metadata: ContextParallelMetadata, rank: int, group: tuple +): + """ + | +-----------before allgather------------+| + | | cp_rank0: block0, block7 | + | | cp_rank1: block1, block6 | + | | cp_rank2: block2, block5 | + | | cp_rank3: block3, block4 | + | + | +----------before rerange---------------+| + | block0 | block7 | block1 | block6 | block2 | block5 | block3 | block4 | + | + | +--------------result-------------------+ + | block0 | block1 | block2 | block3 | block4 | block5 | block6 | block7 | + | +-------------------------+ + """ + x = token_all_gather( + x, + group, + scattered_num_tokens=cp_metadata.per_rank_actual_token, + ) + cp_segment_num = len(cp_metadata.split_list) + inverse_index = list(range(0, cp_segment_num, 2)) + list( + range(cp_segment_num - 1, 0, -2) + ) + x_list = torch.split(x, cp_metadata.inverse_split_list) + output = torch.cat([x_list[i] for i in inverse_index]) + return output diff --git a/python/tokenspeed/runtime/layers/vocab_parallel_embedding.py b/python/tokenspeed/runtime/layers/vocab_parallel_embedding.py new file mode 100755 index 0000000..481ca8c --- /dev/null +++ b/python/tokenspeed/runtime/layers/vocab_parallel_embedding.py @@ -0,0 +1,601 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + + +from collections.abc import Sequence +from dataclasses import dataclass + +import torch +import torch.nn.functional as F +from torch.nn.parameter import Parameter, UninitializedParameter + +from tokenspeed.runtime.distributed.comm_ops import all_reduce +from tokenspeed.runtime.distributed.utils import divide +from tokenspeed.runtime.layers.parameter import BaseWeightParameter +from tokenspeed.runtime.layers.quantization.base_config import ( + QuantizationConfig, + QuantizeMethodBase, + method_has_implemented_embedding, +) +from tokenspeed.runtime.utils import set_weight_attrs + +DEFAULT_VOCAB_PADDING_SIZE = 64 + + +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], + 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: torch.Tensor | None = 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) + + +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 +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.jit.script 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. + + The vocabulary size is padded so 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: torch.dtype | None = None, + org_num_embeddings: int | None = None, + padding_size: int = DEFAULT_VOCAB_PADDING_SIZE, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + tp_rank: int | None = None, + tp_size: int | None = None, + tp_group: tuple[int, ...] | None = None, + use_presharded_weights: bool = False, + ): + super().__init__() + self.quant_config = quant_config + + if tp_rank is None: + assert tp_size is None + assert tp_group is None + self.tp_rank, self.tp_size, self.tp_group = 0, 1, None + else: + assert tp_size is not None + assert tp_group is not None + self.tp_rank, self.tp_size, self.tp_group = tp_rank, tp_size, tp_group + + self.num_embeddings = num_embeddings + self.padding_size = padding_size + self.org_vocab_size = org_num_embeddings or num_embeddings + 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, + self.tp_rank, + self.tp_size, + ) + self.embedding_dim = embedding_dim + + linear_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 + linear_method_implements_embedding = method_has_implemented_embedding( + type(linear_method) + ) + if is_embedding_layer and not linear_method_implements_embedding: + raise NotImplementedError( + f"The class {type(linear_method).__name__} must implement " + "the 'embedding' method, see UnquantizedEmbeddingMethod." + ) + + self.linear_method: QuantizeMethodBase = linear_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.linear_method.create_weights( + self, + self.embedding_dim, + [self.num_embeddings_per_partition], + 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) -> list[int] | None: + """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 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, BaseWeightParameter) + else param.pack_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_: torch.Tensor, reduce_results: bool = True + ) -> torch.Tensor: + """Forward pass for the vocabulary parallel embedding. + + Args: + input_: The input tensor. + reduce_results: Whether to reduce the results across all the model parallel GPUs. + + Returns: + The output tensor. + + """ + 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: + # Single-rank (DP / replicated) path has no shard mask, so an + # out-of-range id (e.g. CUDA-graph capture warmup where the drafter + # feeds argmax over uninitialized logits) hits F.embedding OOB. + # Clamp here for parity with the tp_size>1 mask path; under normal + # inference upstream guarantees id < num_embeddings_padded so this + # is a no-op. + masked_input = input_.clamp(min=0, max=self.num_embeddings_padded - 1) + + # Get the embeddings. + output_parallel = self.linear_method.embedding(self, masked_input) + + # Mask the output embedding. + if self.tp_size > 1: + output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0) + + if reduce_results: + output = all_reduce(output_parallel, self.tp_group) + else: + output = 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}" + 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: torch.dtype | None = None, + org_num_embeddings: int | None = None, + padding_size: int = DEFAULT_VOCAB_PADDING_SIZE, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + tp_rank: int | None = None, + tp_size: int | None = None, + tp_group: tuple[int, ...] | None = None, + use_presharded_weights: bool = False, + ): + super().__init__( + num_embeddings, + embedding_dim, + params_dtype, + org_num_embeddings, + padding_size, + quant_config, + prefix, + tp_rank=tp_rank, + tp_size=tp_size, + tp_group=tp_group, + use_presharded_weights=use_presharded_weights, + ) + self.quant_config = quant_config + 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.""" + 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/tokenspeed/runtime/metrics/collector.py b/python/tokenspeed/runtime/metrics/collector.py new file mode 100755 index 0000000..ef20060 --- /dev/null +++ b/python/tokenspeed/runtime/metrics/collector.py @@ -0,0 +1,535 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from dataclasses import dataclass +from enum import Enum + +from tokenspeed.runtime.utils.env import envs + +TOKENSPEED_TEST_REQUEST_TIME_STATS = envs.TOKENSPEED_TEST_REQUEST_TIME_STATS.get() + + +def _raise_for_negative_durations(**durations: float) -> None: + negative_durations = [ + f"{name}={duration} < 0" for name, duration in durations.items() if duration < 0 + ] + if negative_durations: + raise ValueError(" or ".join(negative_durations)) + + +@dataclass +class TimeStats: + """ + Store the timestamps for each stage of a request. + + Unified: wait_queue -> forward -> completion + Prefill: bootstrap_queue -> wait_queue -> forward -> transfer_queue -> completion + Decode: prealloc_queue -> transfer_queue -> wait_queue -> forward -> completion + """ + + lb_entry_time: float = 0.0 + wait_queue_entry_time: float = 0.0 + forward_entry_time: float = 0.0 + completion_time: float = 0.0 + prefill_bootstrap_queue_entry_time: float = 0.0 + prefill_transfer_queue_entry_time: float = 0.0 + decode_prealloc_queue_entry_time: float = 0.0 + decode_transfer_queue_entry_time: float = 0.0 + + class RequestType(Enum): + UNIFIED = "unified" + PREFILL = "prefill" + DECODE = "decode" + INVALID = "invalid" + + def get_queueing_time(self) -> float: + return self.forward_entry_time - self.wait_queue_entry_time + + def __str__(self) -> str: + # if unified + _type = self.get_type() + + if _type == self.RequestType.UNIFIED: + queue_duration = self.forward_entry_time - self.wait_queue_entry_time + forward_duration = self.completion_time - self.forward_entry_time + + if TOKENSPEED_TEST_REQUEST_TIME_STATS: + _raise_for_negative_durations( + queue_duration=queue_duration, + forward_duration=forward_duration, + ) + + return f"queue_duration={self.format_duration(queue_duration)}, forward_duration={self.format_duration(forward_duration)}, start_time={self.wait_queue_entry_time}" + if _type == self.RequestType.PREFILL: + bootstrap_duration = ( + self.wait_queue_entry_time - self.prefill_bootstrap_queue_entry_time + ) + + queue_duration = self.forward_entry_time - self.wait_queue_entry_time + + forward_duration = self.completion_time - self.forward_entry_time + + if TOKENSPEED_TEST_REQUEST_TIME_STATS: + _raise_for_negative_durations( + bootstrap_duration=bootstrap_duration, + queue_duration=queue_duration, + forward_duration=forward_duration, + ) + return f"bootstrap_duration={self.format_duration(bootstrap_duration)}, queue_duration={self.format_duration(queue_duration)}, forward_duration={self.format_duration(forward_duration)}, start_time={self.prefill_bootstrap_queue_entry_time}" + # if decode + if _type == self.RequestType.DECODE: + prealloc_duration = ( + self.decode_transfer_queue_entry_time + - self.decode_prealloc_queue_entry_time + ) + + transfer_duration = ( + self.wait_queue_entry_time - self.decode_transfer_queue_entry_time + ) + queue_duration = self.forward_entry_time - self.wait_queue_entry_time + forward_duration = self.completion_time - self.forward_entry_time + + if TOKENSPEED_TEST_REQUEST_TIME_STATS: + _raise_for_negative_durations( + prealloc_duration=prealloc_duration, + transfer_duration=transfer_duration, + queue_duration=queue_duration, + forward_duration=forward_duration, + ) + + return f"prealloc_duration={self.format_duration(prealloc_duration)}, transfer_duration={self.format_duration(transfer_duration)}, queue_duration={self.format_duration(queue_duration)}, forward_duration={self.format_duration(forward_duration)}, start_time={self.decode_prealloc_queue_entry_time}" + return "Invalid Time Stats" + + def format_duration(self, duration: float) -> str: + return f"{duration * 1e3:.2f}ms" + + def get_type(self) -> RequestType: + """Determine the type of request based on timestamp values.""" + if ( + self.prefill_bootstrap_queue_entry_time == 0.0 + and self.prefill_transfer_queue_entry_time == 0.0 + and self.decode_prealloc_queue_entry_time == 0.0 + and self.decode_transfer_queue_entry_time == 0.0 + ): + return self.RequestType.UNIFIED + elif ( + self.prefill_bootstrap_queue_entry_time > 0.0 + and self.prefill_transfer_queue_entry_time > 0.0 + ): + return self.RequestType.PREFILL + elif ( + self.decode_prealloc_queue_entry_time > 0.0 + and self.decode_transfer_queue_entry_time > 0.0 + and self.wait_queue_entry_time > 0.0 + ): + return self.RequestType.DECODE + else: + return self.RequestType.INVALID + + +@dataclass +class RequestFinishStats: + prompt_tokens: int + generation_tokens: int + e2e_latency: float + cached_prompt_tokens: int = 0 + finished_ok: bool = True + + +class EngineMetrics: + def __init__( + self, + labels: dict[str, str], + *, + enabled: bool, + registry=None, + ) -> None: + self.enabled = enabled + self.labels = labels + + if self.enabled: + self._init_prometheus(labels, registry=registry) + + def _init_prometheus(self, labels: dict[str, str], *, registry=None) -> None: + from prometheus_client import Counter, Gauge, Histogram + + labelnames = list(labels.keys()) + kw = {"registry": registry} if registry is not None else {} + self.num_requests_running = Gauge( + name="tokenspeed:num_requests_running", + documentation="Requests with scheduler-side generation state (decode path).", + labelnames=labelnames, + multiprocess_mode="livemax", + **kw, + ) + self.num_requests_waiting = Gauge( + name="tokenspeed:num_requests_waiting", + documentation="Requests waiting in the C++ scheduler queue.", + labelnames=labelnames, + multiprocess_mode="livemax", + **kw, + ) + # Wire name follows vLLM's `vllm:kv_cache_usage_perc` for s/vllm:/tokenspeed:/g + # parity even though the value is a 0-1 ratio, not a percentage. + self.kv_cache_usage_ratio = Gauge( + name="tokenspeed:kv_cache_usage_perc", + documentation="Fraction of device KV pages in use (0-1).", + labelnames=labelnames, + multiprocess_mode="livemax", + **kw, + ) + self.iteration_tokens_total = Histogram( + name="tokenspeed:iteration_tokens_total", + documentation="Tokens scheduled in one scheduler forward step.", + labelnames=labelnames, + buckets=[ + 0.0, + 1.0, + 2.0, + 4.0, + 8.0, + 16.0, + 32.0, + 64.0, + 128.0, + 256.0, + 512.0, + 1024.0, + 2048.0, + 4096.0, + 8192.0, + ], + **kw, + ) + self.spec_decode_num_accepted_tokens = Counter( + name="tokenspeed:spec_decode_num_accepted_tokens", + documentation=( + "Accepted speculative draft tokens (excludes the bonus token sampled " + "after verify)." + ), + labelnames=labelnames, + **kw, + ) + self.spec_decode_num_draft_tokens = Counter( + name="tokenspeed:spec_decode_num_draft_tokens", + documentation="Draft tokens proposed across verify steps.", + labelnames=labelnames, + **kw, + ) + self.spec_decode_num_drafts = Counter( + name="tokenspeed:spec_decode_num_drafts", + documentation="Number of speculative verify rounds (per request-slot).", + labelnames=labelnames, + **kw, + ) + self.num_nan_aborted_requests = Counter( + name="tokenspeed:num_nan_aborted_requests", + documentation=( + "Requests terminated by the NaN guard (NaN in logits or an " + "out-of-vocab sampled token id)." + ), + labelnames=labelnames, + **kw, + ) + + def set_scheduler_snapshot( + self, *, running: int, waiting: int, kv_cache_usage_ratio: float + ) -> None: + if not self.enabled: + return + self.num_requests_running.labels(**self.labels).set(running) + self.num_requests_waiting.labels(**self.labels).set(waiting) + self.kv_cache_usage_ratio.labels(**self.labels).set(kv_cache_usage_ratio) + + def observe_iteration_tokens(self, num_tokens: float) -> None: + if self.enabled and num_tokens >= 0: + self.iteration_tokens_total.labels(**self.labels).observe(num_tokens) + + def record_scheduler_iteration( + self, + *, + running: int, + waiting: int, + num_active_pages: int, + num_total_pages: int, + num_iteration_tokens: int, + ) -> None: + if not self.enabled: + return + ratio = num_active_pages / num_total_pages if num_total_pages else 0.0 + self.set_scheduler_snapshot( + running=running, + waiting=waiting, + kv_cache_usage_ratio=ratio, + ) + if num_iteration_tokens > 0: + self.observe_iteration_tokens(float(num_iteration_tokens)) + + def record_spec_decode_step( + self, + *, + num_decode_slots: int, + accepted_draft_tokens: int, + draft_width: int, + ) -> None: + if not self.enabled or num_decode_slots <= 0: + return + self.spec_decode_num_drafts.labels(**self.labels).inc(num_decode_slots) + self.spec_decode_num_draft_tokens.labels(**self.labels).inc( + num_decode_slots * draft_width + ) + self.spec_decode_num_accepted_tokens.labels(**self.labels).inc( + max(0, accepted_draft_tokens) + ) + + def record_nan_abort(self) -> None: + if not self.enabled: + return + self.num_nan_aborted_requests.labels(**self.labels).inc() + + +class RequestMetrics: + def __init__( + self, + labels: dict[str, str], + *, + enabled: bool, + registry=None, + ) -> None: + self.enabled = enabled + self.labels = labels + + if self.enabled: + self._init_prometheus(labels, registry=registry) + + def _init_prometheus(self, labels: dict[str, str], *, registry=None) -> None: + # We need to import prometheus_client after setting the env variable PROMETHEUS_MULTIPROC_DIR + from prometheus_client import Counter, Histogram + + labelnames = list(labels.keys()) + kw = {"registry": registry} if registry is not None else {} + + self.prompt_tokens_total = Counter( + name="tokenspeed:prompt_tokens", + documentation="Number of prefill tokens processed.", + labelnames=labelnames, + **kw, + ) + + self.generation_tokens_total = Counter( + name="tokenspeed:generation_tokens", + documentation="Number of generation tokens processed.", + labelnames=labelnames, + **kw, + ) + + # vLLM has no direct equivalent; tokenspeed-only Counter that tracks + # every finished request regardless of finish reason. + self.num_requests_total = Counter( + name="tokenspeed:num_requests", + documentation="Number of requests processed.", + labelnames=labelnames, + **kw, + ) + + self.request_success_total = Counter( + name="tokenspeed:request_success", + documentation="Requests that finished without an abort-style finish.", + labelnames=labelnames, + **kw, + ) + + self.prefix_cache_hits_total = Counter( + name="tokenspeed:prefix_cache_hits", + documentation=( + "Prompt tokens served from prefix cache. Hit ratio = " + "prefix_cache_hits_total / prompt_tokens_total." + ), + labelnames=labelnames, + **kw, + ) + + self.histogram_time_to_first_token = Histogram( + name="tokenspeed:time_to_first_token_seconds", + documentation="Histogram of time to first token in seconds.", + labelnames=labelnames, + buckets=[ + 0.1, + 0.3, + 0.5, + 0.7, + 0.9, + 1, + 2, + 4, + 6, + 8, + 10, + 20, + 40, + 60, + 80, + 120, + 160, + ], + **kw, + ) + + self.histogram_time_per_output_token = Histogram( + name="tokenspeed:request_time_per_output_token_seconds", + documentation="Histogram of time per output token in seconds.", + labelnames=labelnames, + buckets=[ + 0.002, + 0.005, + 0.010, + 0.020, + 0.030, + 0.040, + 0.050, + 0.060, + 0.070, + 0.080, + 0.090, + 0.100, + 0.150, + 0.200, + 0.300, + 0.400, + 0.600, + 0.800, + 1.000, + 2.000, + ], + **kw, + ) + + self.histogram_inter_token_latency_seconds = Histogram( + name="tokenspeed:inter_token_latency_seconds", + documentation="Histogram of inter-token latency in seconds.", + labelnames=labelnames, + buckets=[ + 0.002, + 0.004, + 0.006, + 0.008, + 0.010, + 0.015, + 0.020, + 0.025, + 0.030, + 0.035, + 0.040, + 0.050, + 0.075, + 0.100, + 0.150, + 0.200, + 0.300, + 0.400, + 0.500, + 0.750, + 1.000, + 2.000, + ], + **kw, + ) + + self.histogram_e2e_request_latency = Histogram( + name="tokenspeed:e2e_request_latency_seconds", + documentation="Histogram of End-to-end request latency in seconds", + labelnames=labelnames, + buckets=[ + 0.1, + 0.2, + 0.4, + 0.8, + 1, + 2, + 5, + 10, + 20, + 40, + 60, + 80, + 100, + 150, + 200, + 250, + 300, + 350, + 500, + 1000, + ], + **kw, + ) + + def record_request_finish(self, stats: RequestFinishStats) -> None: + if not self.enabled: + return + self.prompt_tokens_total.labels(**self.labels).inc(stats.prompt_tokens) + self.generation_tokens_total.labels(**self.labels).inc(stats.generation_tokens) + self.num_requests_total.labels(**self.labels).inc(1) + self.prefix_cache_hits_total.labels(**self.labels).inc( + stats.cached_prompt_tokens + ) + if stats.finished_ok: + self.request_success_total.labels(**self.labels).inc(1) + self.histogram_e2e_request_latency.labels(**self.labels).observe( + stats.e2e_latency + ) + if stats.generation_tokens >= 1: + self.histogram_time_per_output_token.labels(**self.labels).observe( + stats.e2e_latency / stats.generation_tokens + ) + + def observe_time_to_first_token(self, value: float) -> None: + if not self.enabled: + return + self.histogram_time_to_first_token.labels(**self.labels).observe(value) + + def observe_inter_token_latency(self, interval: float, num_new_tokens: int) -> None: + if not self.enabled: + return + adjusted_interval = interval / num_new_tokens + # A faster version of the Histogram::observe which observes multiple values at the same time. + # reference: https://github.com/prometheus/client_python/blob/v0.21.1/prometheus_client/metrics.py#L639 + his = self.histogram_inter_token_latency_seconds.labels(**self.labels) + his._sum.inc(interval) + + for i, bound in enumerate(his._upper_bounds): + if adjusted_interval <= bound: + his._buckets[i].inc(num_new_tokens) + break + + +class KVTransferMetrics: + + def __init__(self, labels: dict[str, str], metrics_reporters: list[str]) -> None: + pass + + def record_kv_transfer_timeout(self) -> None: + return + + def record_kv_transfer_failure(self) -> None: + return + + def observe_kv_transfer_latency(self, transfer_time_seconds: float) -> None: + return diff --git a/python/tokenspeed/runtime/metrics/func_timer.py b/python/tokenspeed/runtime/metrics/func_timer.py new file mode 100755 index 0000000..79a85af --- /dev/null +++ b/python/tokenspeed/runtime/metrics/func_timer.py @@ -0,0 +1,45 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Record the latency of selected functions.""" + +from tokenspeed.runtime.metrics.utils import exponential_buckets + +enable_metrics = False + + +def enable_func_timer(): + """Enable Prometheus-backed function latency metrics.""" + # We need to import prometheus_client after setting the env variable `PROMETHEUS_MULTIPROC_DIR` + from prometheus_client import Histogram + + global enable_metrics, FUNC_LATENCY + enable_metrics = True + + FUNC_LATENCY = Histogram( + "tokenspeed:func_latency_seconds", + "Function latency in seconds", + # captures latency in range [50ms - ~50s] + buckets=exponential_buckets(start=0.05, width=1.5, length=18), + labelnames=["name"], + ) + + +FUNC_LATENCY = None diff --git a/python/tokenspeed/runtime/metrics/utils.py b/python/tokenspeed/runtime/metrics/utils.py new file mode 100755 index 0000000..463fae2 --- /dev/null +++ b/python/tokenspeed/runtime/metrics/utils.py @@ -0,0 +1,28 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Utilities for Prometheus Metrics.""" + + +def exponential_buckets(start: float, width: float, length: int) -> list[float]: + buckets = [] + for i in range(length): + buckets.append(start * (width**i)) + return buckets diff --git a/python/tokenspeed/runtime/model_loader/__init__.py b/python/tokenspeed/runtime/model_loader/__init__.py new file mode 100755 index 0000000..d780fda --- /dev/null +++ b/python/tokenspeed/runtime/model_loader/__init__.py @@ -0,0 +1,49 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from torch import nn + +from tokenspeed.runtime.configs.device_config import DeviceConfig +from tokenspeed.runtime.configs.load_config import LoadConfig +from tokenspeed.runtime.configs.model_config import ModelConfig +from tokenspeed.runtime.model_loader.loader import BaseModelLoader, get_model_loader +from tokenspeed.runtime.model_loader.utils import get_model_architecture + + +def get_model( + *, + model_config: ModelConfig, + load_config: LoadConfig, + device_config: DeviceConfig, +) -> nn.Module: + """Load and return a model for the given runtime configuration.""" + loader = get_model_loader(load_config) + return loader.load_model( + model_config=model_config, + device_config=device_config, + ) + + +__all__ = [ + "get_model", + "get_model_loader", + "BaseModelLoader", + "get_model_architecture", +] diff --git a/python/tokenspeed/runtime/model_loader/loader.py b/python/tokenspeed/runtime/model_loader/loader.py new file mode 100755 index 0000000..2716089 --- /dev/null +++ b/python/tokenspeed/runtime/model_loader/loader.py @@ -0,0 +1,636 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + + +# ruff: noqa: SIM117 +import collections +import dataclasses +import glob +import os +from abc import ABC, abstractmethod +from collections.abc import Generator, Iterable, Iterator +from contextlib import contextmanager +from typing import Any, cast + +import huggingface_hub +import torch +import yaml +from tokenspeed_kernel.platform import current_platform +from torch import nn +from transformers.utils import SAFE_WEIGHTS_INDEX_NAME + +from tokenspeed.runtime.configs.device_config import DeviceConfig +from tokenspeed.runtime.configs.load_config import LoadConfig, LoadFormat +from tokenspeed.runtime.configs.model_config import ModelConfig +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.model_loader.utils import ( + get_model_architecture, + set_default_torch_dtype, +) +from tokenspeed.runtime.model_loader.weight_utils import ( + download_safetensors_index_file_from_hf, + download_weights_from_hf, + filter_duplicate_safetensors_files, + filter_files_not_needed_for_inference, + get_quant_config, + initialize_dummy_weights, + np_cache_weights_iterator, + pt_weights_iterator, + safetensors_weights_iterator, +) +from tokenspeed.runtime.models.extensible import ExtensibleLM +from tokenspeed.runtime.utils import get_colorful_logger, is_pin_memory_available + + +@contextmanager +def device_loading_context( + module: torch.nn.Module, target_device: torch.device +) -> Iterator[torch.nn.Module]: + if target_device.type == "cpu": + # If target is CPU, no need to move anything + yield module + return + + original_device_states: dict[str, torch.device] = {} + + # Store original device states and move parameters to GPU if they're on CPU + for name, p in module.named_parameters(): + if p.device.type == "cpu": + original_device_states[name] = p.device + p.data = p.data.to(target_device) + # Parameters already on target device are not touched + + try: + yield module + + finally: + # Restore parameters to their original devices, ignoring new parameters + pin_memory = is_pin_memory_available() + for name, p in module.named_parameters(): + if name in original_device_states: + original_device: torch.device = original_device_states[name] + if original_device.type == "cpu": + # `torch.empty_like` does not support `pin_memory` argument + cpu_data = torch.empty_strided( + size=p.data.size(), + stride=p.data.stride(), + dtype=p.data.dtype, + layout=p.data.layout, + device="cpu", + pin_memory=pin_memory, + ) + cpu_data.copy_(p.data) + p.data = cpu_data + else: + p.data = p.data.to(original_device) + # New parameters or parameters already on target device are untouched + + +logger = get_colorful_logger(__name__) + + +def _get_quantization_config( + model_config: ModelConfig, load_config: LoadConfig +) -> QuantizationConfig | None: + """Get the quantization config.""" + if model_config.quantization is not None: + quant_config = get_quant_config(model_config, load_config) + platform = current_platform() + capability = platform.arch_version.major * 10 + platform.arch_version.minor + if capability < quant_config.get_min_capability(): + raise ValueError( + f"The quantization method {model_config.quantization} " + "is not supported for the current GPU. " + f"Minimum capability: {quant_config.get_min_capability()}. " + f"Current capability: {capability}." + ) + supported_dtypes = quant_config.get_supported_act_dtypes() + if model_config.dtype not in supported_dtypes: + raise ValueError( + f"{model_config.dtype} is not supported for quantization " + f"method {model_config.quantization}. Supported dtypes: " + f"{supported_dtypes}" + ) + return quant_config + return None + + +def _initialize_model( + model_config: ModelConfig, + load_config: LoadConfig, +) -> nn.Module: + """Initialize a model with the given configurations.""" + model_class, _ = get_model_architecture(model_config) + quant_config = _get_quantization_config(model_config, load_config) + mapping = model_config.mapping + # Only VLM wrappers accept these kwargs. + extra_kwargs: dict = {} + if model_config.is_multimodal: + extra_kwargs["is_multimodal_active"] = model_config.is_multimodal_active + extra_kwargs["mm_attention_backend"] = model_config.mm_attention_backend + return model_class( + config=model_config.hf_config, + mapping=mapping, + quant_config=quant_config, + **extra_kwargs, + ) + + +class BaseModelLoader(ABC): + """Base class for model loaders.""" + + def __init__(self, load_config: LoadConfig) -> None: + self.load_config = load_config + + @abstractmethod + def download_model(self, model_config: ModelConfig) -> None: + """Download a model so that it can be immediately loaded.""" + raise NotImplementedError + + @abstractmethod + def load_model( + self, + *, + model_config: ModelConfig, + device_config: DeviceConfig, + ) -> nn.Module: + """Load a model with the given configurations.""" + raise NotImplementedError + + +class DefaultModelLoader(BaseModelLoader): + """Model loader that can load different file types from disk.""" + + @dataclasses.dataclass + class Source: + """A source for weights.""" + + model_or_path: str + """The model ID or path.""" + + revision: str | None + """The optional model revision.""" + + prefix: str = "" + """A prefix to prepend to all weights.""" + + fall_back_to_pt: bool = True + """Whether .pt weights can be used.""" + + def __init__(self, load_config: LoadConfig) -> None: + super().__init__(load_config) + if load_config.model_loader_extra_config: + raise ValueError( + f"Model loader extra config is not supported for " + f"load format {load_config.load_format}" + ) + + def _maybe_download_from_modelscope( + self, model: str, revision: str | None + ) -> str | None: + """Download model from ModelScope hub if TOKENSPEED_USE_MODELSCOPE is True. + + Returns the path to the downloaded model, or None if the model is not + downloaded from ModelScope.""" + from tokenspeed.runtime.utils.env import envs + + if envs.TOKENSPEED_USE_MODELSCOPE.is_set(): + # download model from ModelScope hub, + # lazy import so that modelscope is not required for normal use. + from modelscope.hub.snapshot_download import snapshot_download + + if not os.path.exists(model): + model_path = snapshot_download( + model_id=model, + cache_dir=self.load_config.download_dir, + local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, + revision=revision, + ignore_file_pattern=self.load_config.ignore_patterns, + ) + else: + model_path = model + return model_path + return None + + def _prepare_weights( + self, model_name_or_path: str, revision: str | None, fall_back_to_pt: bool + ) -> tuple[str, list[str], bool]: + """Prepare weights for the model. + + If the model is not local, it will be downloaded.""" + model_name_or_path = ( + self._maybe_download_from_modelscope(model_name_or_path, revision) + or model_name_or_path + ) + + is_local = os.path.isdir(model_name_or_path) + load_format = self.load_config.load_format + use_safetensors = False + index_file = SAFE_WEIGHTS_INDEX_NAME + # Some quantized models use .pt files for storing the weights. + if load_format == LoadFormat.AUTO: + allow_patterns = ["*.safetensors", "*.bin"] + elif load_format == LoadFormat.SAFETENSORS: + use_safetensors = True + allow_patterns = ["*.safetensors"] + elif load_format == LoadFormat.MISTRAL: + use_safetensors = True + allow_patterns = ["consolidated*.safetensors"] + index_file = "consolidated.safetensors.index.json" + elif load_format == LoadFormat.PT: + allow_patterns = ["*.pt"] + elif load_format == LoadFormat.NPCACHE: + allow_patterns = ["*.bin"] + else: + raise ValueError(f"Unknown load_format: {load_format}") + + if fall_back_to_pt: + allow_patterns += ["*.pt"] + + if not is_local: + hf_folder = download_weights_from_hf( + model_name_or_path, + self.load_config.download_dir, + allow_patterns, + revision, + ignore_patterns=self.load_config.ignore_patterns, + ) + else: + hf_folder = model_name_or_path + + hf_weights_files: list[str] = [] + for pattern in allow_patterns: + hf_weights_files += glob.glob(os.path.join(hf_folder, pattern)) + if len(hf_weights_files) > 0: + if pattern == "*.safetensors": + use_safetensors = True + break + + if use_safetensors: + # For models like Mistral-7B-Instruct-v0.3 + # there are both sharded safetensors files and a consolidated + # safetensors file. Using both breaks. + # Here, we download the `model.safetensors.index.json` and filter + # any files not found in the index. + if not is_local: + download_safetensors_index_file_from_hf( + model_name_or_path, + index_file, + self.load_config.download_dir, + revision, + ) + hf_weights_files = filter_duplicate_safetensors_files( + hf_weights_files, hf_folder, index_file + ) + else: + hf_weights_files = filter_files_not_needed_for_inference(hf_weights_files) + + if len(hf_weights_files) == 0: + raise RuntimeError( + f"Cannot find any model weights with `{model_name_or_path}`" + ) + + return hf_folder, hf_weights_files, use_safetensors + + def _get_weights_iterator( + self, source: "Source" + ) -> Generator[tuple[str, torch.Tensor], None, None]: + """Get an iterator for the model weights based on the load format.""" + hf_folder, hf_weights_files, use_safetensors = self._prepare_weights( + source.model_or_path, source.revision, source.fall_back_to_pt + ) + if self.load_config.load_format == LoadFormat.NPCACHE: + # Currently np_cache only support *.bin checkpoints + if use_safetensors: + raise ValueError("np_cache only supports PyTorch checkpoint shards.") + weights_iterator = np_cache_weights_iterator( + source.model_or_path, + self.load_config.download_dir, + hf_folder, + hf_weights_files, + ) + elif use_safetensors: + weights_iterator = safetensors_weights_iterator( + hf_weights_files, + prefetch=self.load_config.weight_loader_prefetch_checkpoints, + prefetch_num_threads=self.load_config.weight_loader_prefetch_num_threads, + ) + else: + weights_iterator = pt_weights_iterator(hf_weights_files) + + # Apply the prefix. + return ((source.prefix + name, tensor) for (name, tensor) in weights_iterator) + + def _get_all_weights( + self, + model_config: ModelConfig, + model: nn.Module, + ) -> Generator[tuple[str, torch.Tensor], None, None]: + + primary_weights = DefaultModelLoader.Source( + model_config.model_path, + model_config.revision, + prefix="", + fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", False), + ) + yield from self._get_weights_iterator(primary_weights) + + secondary_weights = cast( + Iterable[DefaultModelLoader.Source], getattr(model, "secondary_weights", ()) + ) + for source in secondary_weights: + yield from self._get_weights_iterator(source) + + def download_model(self, model_config: ModelConfig) -> None: + self._prepare_weights( + model_config.model_path, model_config.revision, fall_back_to_pt=True + ) + + def load_model( + self, + *, + model_config: ModelConfig, + device_config: DeviceConfig, + ) -> nn.Module: + target_device = torch.device(device_config.device) + with set_default_torch_dtype(model_config.dtype): + with target_device: + model = _initialize_model( + model_config, + self.load_config, + ) + + model.load_weights(self._get_all_weights(model_config, model)) + + for _, module in model.named_modules(): + quant_method = getattr(module, "quant_method", None) + if quant_method is not None: + # When quant methods need to process weights after loading + # (for repacking, quantizing, etc), they expect parameters + # to be on the global target device. This scope is for the + # case where cpu offloading is used, where we will move the + # parameters onto device for processing and back off after. + with device_loading_context(module, target_device): + quant_method.process_weights_after_loading(module) + + process_method = getattr(module, "process_weights_after_loading", None) + if process_method is not None: + with device_loading_context(module, target_device): + module.process_weights_after_loading(module) + + post_quant_warmup = getattr(model, "post_quant_warmup", None) + if callable(post_quant_warmup): + post_quant_warmup() + + return model.eval() + + +class ExtensibleModelLoader: + def __init__(self, load_config: LoadConfig) -> None: + load_config.load_format = LoadFormat.AUTO + self.base_lm_loader = DefaultModelLoader(load_config) + self.ext_yaml = load_config.ext_yaml + + def load_model( + self, + model_config: ModelConfig, + device_config: DeviceConfig, + ) -> nn.Module: + with open(self.ext_yaml) as f: + ext_config = yaml.safe_load(f) + base_lm = self.base_lm_loader.load_model( + model_config=model_config, device_config=device_config + ) + ext_lm = ExtensibleLM(base_lm=base_lm, ext_config=ext_config) + return ext_lm + + +class DummyModelLoader(BaseModelLoader): + """Model loader that will set model weights to random values.""" + + def __init__(self, load_config: LoadConfig) -> None: + super().__init__(load_config) + if load_config.model_loader_extra_config: + raise ValueError( + f"Model loader extra config is not supported for " + f"load format {load_config.load_format}" + ) + + def download_model(self, model_config: ModelConfig) -> None: + pass # Nothing to download + + def load_model( + self, + *, + model_config: ModelConfig, + device_config: DeviceConfig, + ) -> nn.Module: + with set_default_torch_dtype(model_config.dtype): + with torch.device(device_config.device): + model = _initialize_model( + model_config, + self.load_config, + ) + if getattr(model, "post_load_weights", None): + model.post_load_weights() + + for _, module in model.named_modules(): + quant_method = getattr(module, "quant_method", None) + if quant_method is not None: + quant_method.process_weights_after_loading(module) + process_method = getattr(module, "process_weights_after_loading", None) + if process_method is not None: + module.process_weights_after_loading(module) + + post_quant_warmup = getattr(model, "post_quant_warmup", None) + if callable(post_quant_warmup): + post_quant_warmup() + + # For accurate performance evaluation, we assign + # random values to the weights. + initialize_dummy_weights(model) + return model.eval() + + +class ShardedStateLoader(BaseModelLoader): + """ + Model loader that directly loads each worker's model state dict, 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. See + `examples/save_sharded_state.py` for creating a sharded checkpoint. + """ + + DEFAULT_PATTERN = "model-rank-{rank}-part-{part}.safetensors" + + def __init__(self, load_config: LoadConfig): + super().__init__(load_config) + extra_config = ( + {} + if load_config.model_loader_extra_config is None + else load_config.model_loader_extra_config.copy() + ) + self.pattern = extra_config.pop("pattern", self.DEFAULT_PATTERN) + if extra_config: + raise ValueError( + f"Unexpected extra config keys for load format " + f"{load_config.load_format}: " + f"{load_config.model_loader_extra_config.keys()}" + ) + + @staticmethod + def _filter_subtensors(tensors: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: + """ + Filter out all tensors that share the same memory or a subset of the + memory of another tensor. + """ + same_storage_groups: dict[Any, list[tuple[str, torch.Tensor]]] = ( + collections.defaultdict(list) + ) + for key, tensor in tensors.items(): + if tensor.numel(): + ptr = tensor.untyped_storage().data_ptr() + same_storage_groups[tensor.device, ptr].append((key, tensor)) + + def get_end_ptr(tensor: torch.Tensor) -> int: + return tensor.view(-1)[-1].data_ptr() + tensor.element_size() + + result: dict[str, torch.Tensor] = {} + for group in same_storage_groups.values(): + for k, t in group: + a, b = t.data_ptr(), get_end_ptr(t) + for k2, t2 in group: + if not t2.is_contiguous(): + continue + a2, b2 = t2.data_ptr(), get_end_ptr(t2) + if a < a2 or b2 < b: + continue + if a2 < a or b < b2 or not t.is_contiguous(): + break # t2 covers strictly more memory than t. + if k2 < k: + # Same tensors, keep the one with the smaller key. + break + else: + result[k] = t + return result + + def _prepare_weights(self, model_name_or_path: str, revision: str | None): + if os.path.isdir(model_name_or_path): + return model_name_or_path + + allow_patterns = ["*.safetensors"] + return download_weights_from_hf( + model_name_or_path, + self.load_config.download_dir, + allow_patterns, + revision, + ignore_patterns=self.load_config.ignore_patterns, + ) + + def download_model(self, model_config: ModelConfig) -> None: + self._prepare_weights(model_config.model_path, model_config.revision) + + def load_model( + self, + *, + model_config: ModelConfig, + device_config: DeviceConfig, + ) -> nn.Module: + from safetensors.torch import safe_open + + local_model_path = self._prepare_weights( + model_config.model_path, model_config.revision + ) + + with set_default_torch_dtype(model_config.dtype): + with torch.device(device_config.device): + model = _initialize_model(model_config, self.load_config) + for _, module in model.named_modules(): + quant_method = getattr(module, "quant_method", None) + if quant_method is not None: + quant_method.process_weights_after_loading(module) + process_method = getattr( + module, "process_weights_after_loading", None + ) + if process_method is not None: + module.process_weights_after_loading(module) + rank = model_config.mapping.rank + pattern = os.path.join( + local_model_path, + self.pattern.format(rank=rank, part="*"), + ) + filepaths = glob.glob(pattern) + if not filepaths: + raise ValueError( + f"Could not find checkpoint files '{pattern}', only " + f"pre-sharded checkpoints are currently supported!" + ) + state_dict = self._filter_subtensors(model.state_dict()) + for path in filepaths: + with safe_open(path, framework="pt") as f: + for key in f.keys(): # noqa: SIM118 + tensor = f.get_tensor(key) + # If loading with LoRA enabled, additional padding may + # be added to certain parameters. We only load into a + # narrowed view of the parameter data. + param_data = state_dict[key].data + param_shape = state_dict[key].shape + for dim, size in enumerate(tensor.shape): + if size < param_shape[dim]: + param_data = param_data.narrow(dim, 0, size) + if tensor.shape != param_shape: + logger.warning( + "loading tensor of shape %s into " + "parameter '%s' of shape %s", + tensor.shape, + key, + param_shape, + ) + param_data.copy_(tensor) + state_dict.pop(key) + if state_dict: + raise ValueError(f"Missing keys {tuple(state_dict)} in loaded state!") + + post_quant_warmup = getattr(model, "post_quant_warmup", None) + if callable(post_quant_warmup): + post_quant_warmup() + + return model.eval() + + +def get_model_loader(load_config: LoadConfig) -> BaseModelLoader: + """Get a model loader based on the load format.""" + + if isinstance(load_config.load_format, type): + return load_config.load_format(load_config) + + if load_config.load_format == LoadFormat.DUMMY: + return DummyModelLoader(load_config) + + if load_config.load_format == LoadFormat.SHARDED_STATE: + return ShardedStateLoader(load_config) + + if load_config.load_format == LoadFormat.EXTENSIBLE: + return ExtensibleModelLoader(load_config) + + return DefaultModelLoader(load_config) diff --git a/python/tokenspeed/runtime/model_loader/utils.py b/python/tokenspeed/runtime/model_loader/utils.py new file mode 100755 index 0000000..95ce0a2 --- /dev/null +++ b/python/tokenspeed/runtime/model_loader/utils.py @@ -0,0 +1,60 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + + +"""Utilities for selecting and loading models.""" + +import contextlib +from collections.abc import Iterator + +import torch +from torch import nn + +from tokenspeed.runtime.configs.model_config import ModelConfig +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +@contextlib.contextmanager +def set_default_torch_dtype(dtype: torch.dtype) -> Iterator[None]: + """Sets the default torch dtype to the given dtype.""" + old_dtype = torch.get_default_dtype() + torch.set_default_dtype(dtype) + yield + torch.set_default_dtype(old_dtype) + + +def get_model_architecture(model_config: ModelConfig) -> tuple[type[nn.Module], str]: + from tokenspeed.runtime.models.registry import ModelRegistry + + architectures = getattr(model_config.hf_config, "architectures", []) + # Mixtral only supports the quantization backends listed here in the + # current model registry and loader stack. + mixtral_supported = ["fp8", "compressed-tensors"] + + if ( + model_config.quantization is not None + and model_config.quantization not in mixtral_supported + and "MixtralForCausalLM" in architectures + ): + architectures = ["QuantMixtralForCausalLM"] + + return ModelRegistry.resolve_model_cls(architectures) diff --git a/python/tokenspeed/runtime/model_loader/weight_utils.py b/python/tokenspeed/runtime/model_loader/weight_utils.py new file mode 100755 index 0000000..3f747cf --- /dev/null +++ b/python/tokenspeed/runtime/model_loader/weight_utils.py @@ -0,0 +1,738 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + + +"""Utilities for downloading and initializing model weights.""" + +import concurrent.futures +import fnmatch +import glob +import hashlib +import importlib.util +import json +import os +import tempfile +import threading +import time +from collections.abc import Callable, Generator, Iterable +from typing import ( + Any, +) + +import filelock +import huggingface_hub.constants +import numpy as np +import safetensors.torch +import torch +from huggingface_hub import HfFileSystem, hf_hub_download, snapshot_download +from pydantic import BaseModel, ConfigDict, ValidationInfo, model_validator +from tqdm.auto import tqdm + +from tokenspeed.runtime.configs.load_config import LoadConfig +from tokenspeed.runtime.configs.model_config import ModelConfig +from tokenspeed.runtime.layers.quantization import ( + QuantizationConfig, + get_quantization_config, +) +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) +_PREFETCH_BLOCK_SIZE = 16 * 1024 * 1024 + +# use system-level temp directory for file locks, so that multiple users +# can share the same lock without error. +# lock files in the temp directory will be automatically deleted when the +# system reboots, so users will not complain about annoying lock files +temp_dir = tempfile.gettempdir() + + +def enable_hf_transfer(): + """automatically activates hf_transfer""" + if "HF_HUB_ENABLE_HF_TRANSFER" not in os.environ: + if importlib.util.find_spec("hf_transfer") is not None: + huggingface_hub.constants.HF_HUB_ENABLE_HF_TRANSFER = True + + +enable_hf_transfer() + + +class DisabledTqdm(tqdm): + def __init__(self, *args: Any, **kwargs: Any) -> None: + kwargs["disable"] = True + super().__init__(*args, **kwargs) + + +def get_lock( + model_name_or_path: str, cache_dir: str | None = None +) -> filelock.FileLock: + lock_dir = cache_dir or temp_dir + os.makedirs(os.path.dirname(lock_dir), exist_ok=True) + model_name = model_name_or_path.replace("/", "-") + hash_name = hashlib.sha256(model_name.encode()).hexdigest() + # add hash to avoid conflict with old users' lock files + lock_file_name = hash_name + model_name + ".lock" + # mode 0o666 is required for the filelock to be shared across users + return filelock.FileLock(os.path.join(lock_dir, lock_file_name), mode=0o666) + + +def get_quant_config( + model_config: ModelConfig, load_config: LoadConfig +) -> QuantizationConfig: + quant_cls = get_quantization_config(model_config.quantization) + + # Read the quantization config from the HF model config, if available. + hf_quant_config = getattr(model_config.hf_config, "quantization_config", None) + # some vision model may keep quantization_config in their text_config + hf_text_config = getattr(model_config.hf_config, "text_config", None) + if hf_quant_config is None and hf_text_config is not None: + hf_quant_config = getattr(hf_text_config, "quantization_config", None) + if hf_quant_config is None: + # compressed-tensors uses a compressions_config + hf_quant_config = getattr(model_config.hf_config, "compression_config", None) + if hf_quant_config is not None: + return quant_cls.from_config(hf_quant_config) + model_name_or_path = model_config.model_path + is_local = os.path.isdir(model_name_or_path) + if not is_local: + # Download the config files. + with get_lock(model_name_or_path, load_config.download_dir): + hf_folder = snapshot_download( + model_name_or_path, + revision=model_config.revision, + allow_patterns="*.json", + cache_dir=load_config.download_dir, + local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, + tqdm_class=DisabledTqdm, + ) + else: + hf_folder = model_name_or_path + + possible_config_filenames = quant_cls.get_config_filenames() + + # If the quantization config is not found, use the default config. + if not possible_config_filenames: + return quant_cls() + + config_files = glob.glob(os.path.join(hf_folder, "*.json")) + + quant_config_files = [ + f for f in config_files if any(f.endswith(x) for x in possible_config_filenames) + ] + if len(quant_config_files) == 0: + raise ValueError(f"Cannot find the config file for {model_config.quantization}") + if len(quant_config_files) > 1: + raise ValueError( + f"Found multiple config files for {model_config.quantization}: " + f"{quant_config_files}" + ) + + quant_config_file = quant_config_files[0] + with open(quant_config_file) as f: + config = json.load(f) + + if model_config.quantization == "nvfp4": + if config["producer"]["name"] == "modelopt": + return quant_cls.from_config(config) + else: + raise ValueError( + f"Unsupported quantization config" + f" found for {model_config.quantization} in {f}." + ) + + return quant_cls.from_config(config) + + +def download_weights_from_hf( + model_name_or_path: str, + cache_dir: str | None, + allow_patterns: list[str], + revision: str | None = None, + ignore_patterns: str | list[str] | None = None, +) -> str: + """Download model weights from Hugging Face Hub. + + Args: + model_name_or_path (str): The model name or path. + cache_dir (Optional[str]): The cache directory to store the model + weights. If None, will use HF defaults. + allow_patterns (List[str]): The allowed patterns for the + weight files. Files matched by any of the patterns will be + downloaded. + revision (Optional[str]): The revision of the model. + ignore_patterns (Optional[Union[str, List[str]]]): The patterns to + filter out the weight files. Files matched by any of the patterns + will be ignored. + + Returns: + str: The path to the downloaded model weights. + """ + if not huggingface_hub.constants.HF_HUB_OFFLINE: + # Before we download we look at that is available: + fs = HfFileSystem() + file_list = fs.ls(model_name_or_path, detail=False, revision=revision) + + # depending on what is available we download different things + for pattern in allow_patterns: + matching = fnmatch.filter(file_list, pattern) + if len(matching) > 0: + allow_patterns = [pattern] + break + + logger.info("Using model weights format %s", allow_patterns) + # Use file lock to prevent multiple processes from + # downloading the same model weights at the same time. + with get_lock(model_name_or_path, cache_dir): + hf_folder = snapshot_download( + model_name_or_path, + allow_patterns=allow_patterns, + ignore_patterns=ignore_patterns, + cache_dir=cache_dir, + tqdm_class=DisabledTqdm, + revision=revision, + local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, + ) + return hf_folder + + +def download_safetensors_index_file_from_hf( + model_name_or_path: str, + index_file: str, + cache_dir: str | None, + revision: str | None = None, +) -> None: + """Download hf safetensors index file from Hugging Face Hub. + + Args: + model_name_or_path (str): The model name or path. + cache_dir (Optional[str]): The cache directory to store the model + weights. If None, will use HF defaults. + revision (Optional[str]): The revision of the model. + """ + # Use file lock to prevent multiple processes from + # downloading the same model weights at the same time. + with get_lock(model_name_or_path, cache_dir): + try: + # Download the safetensors index file. + hf_hub_download( + repo_id=model_name_or_path, + filename=index_file, + cache_dir=cache_dir, + revision=revision, + local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, + ) + # If file not found on remote or locally, we should not fail since + # only some models will have index_file. + except huggingface_hub.utils.EntryNotFoundError: + logger.info("No %s found in remote.", index_file) + except huggingface_hub.utils.LocalEntryNotFoundError: + logger.info("No %s found in local cache.", index_file) + + +# For models like Mistral-7B-v0.3, there are both sharded +# safetensors files and a consolidated safetensors file. +# Passing both of these to the weight loader functionality breaks. +# So, we use the index_file to +# look up which safetensors files should be used. +def filter_duplicate_safetensors_files( + hf_weights_files: list[str], hf_folder: str, index_file: str +) -> list[str]: + # model.safetensors.index.json is a mapping from keys in the + # torch state_dict to safetensors file holding that weight. + index_file_name = os.path.join(hf_folder, index_file) + if not os.path.isfile(index_file_name): + return hf_weights_files + + # Iterate through the weight_map (weight_name: safetensors files) + # to identify weights that we should use. + with open(index_file_name) as f: + weight_map = json.load(f)["weight_map"] + weight_files_in_index = set() + for weight_name in weight_map: + weight_files_in_index.add(os.path.join(hf_folder, weight_map[weight_name])) + # Filter out any fields that are not found in the index file. + hf_weights_files = [f for f in hf_weights_files if f in weight_files_in_index] + return hf_weights_files + + +def filter_files_not_needed_for_inference(hf_weights_files: list[str]) -> list[str]: + """ + Exclude files that are not needed for inference. + + See https://github.com/huggingface/transformers/blob/v4.34.0/src/transformers/trainer.py#L227-L233 + """ + blacklist = [ + "training_args.bin", + "optimizer.bin", + "optimizer.pt", + "scheduler.pt", + "scaler.pt", + ] + hf_weights_files = [ + f for f in hf_weights_files if not any(f.endswith(x) for x in blacklist) + ] + return hf_weights_files + + +# explicitly use pure text format, with a newline at the end +# this makes it impossible to see the animation in the progress bar +# but will avoid messing up with ray or multiprocessing, which wraps +# each line of output with some prefix. +_BAR_FORMAT = "{desc}: {percentage:3.0f}% Completed | {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}]\n" # noqa: E501 + + +def np_cache_weights_iterator( + model_name_or_path: str, + cache_dir: str | None, + hf_folder: str, + hf_weights_files: list[str], +) -> Generator[tuple[str, torch.Tensor], None, None]: + """Iterate over the weights in the model np files. + + Will dump the model weights to numpy files if they are not already dumped. + """ + enable_tqdm = ( + not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 + ) + # Convert the model weights from torch tensors to numpy arrays for + # faster loading. + np_folder = os.path.join(hf_folder, "np") + os.makedirs(np_folder, exist_ok=True) + weight_names_file = os.path.join(np_folder, "weight_names.json") + # Use file lock to prevent multiple processes from + # dumping the same model weights to numpy at the same time. + with get_lock(model_name_or_path, cache_dir): + if not os.path.exists(weight_names_file): + weight_names: list[str] = [] + for bin_file in tqdm( + hf_weights_files, + desc="Loading np_cache checkpoint shards", + disable=not enable_tqdm, + bar_format=_BAR_FORMAT, + ): + state = torch.load(bin_file, map_location="cpu") + for name, param in state.items(): + param_path = os.path.join(np_folder, name) + with open(param_path, "wb") as f: + np.save(f, param.cpu().detach().numpy()) + weight_names.append(name) + with open(weight_names_file, "w") as f: + json.dump(weight_names, f) + + with open(weight_names_file) as f: + weight_names = json.load(f) + + for name in weight_names: + param_path = os.path.join(np_folder, name) + with open(param_path, "rb") as f: + param = np.load(f) + yield name, torch.from_numpy(param) + + +def decrypt(fn, key): + raise NotImplementedError() + + +def safetensors_encrypted_weights_iterator( + hf_weights_files: list[str], + is_all_weights_sharded: bool = False, + decryption_key: str | None = None, +): + raise NotImplementedError() + + +def _get_checkpoint_prefetch_rank_info() -> tuple[int, int]: + local_rank = os.getenv("LOCAL_RANK") + local_world_size = os.getenv("LOCAL_WORLD_SIZE") + if local_rank is not None and local_world_size is not None: + try: + rank = int(local_rank) + world_size = int(local_world_size) + if 0 <= rank < world_size: + return rank, world_size + except ValueError: + logger.warning( + "Ignoring invalid LOCAL_RANK/LOCAL_WORLD_SIZE for checkpoint prefetch: " + "%s/%s", + local_rank, + local_world_size, + ) + + if torch.distributed.is_available() and torch.distributed.is_initialized(): + return torch.distributed.get_rank(), torch.distributed.get_world_size() + + return 0, 1 + + +def _prefetch_checkpoint_file(file_path: str) -> int: + bytes_read = 0 + with open(file_path, "rb") as f: + while True: + data = f.read(_PREFETCH_BLOCK_SIZE) + if not data: + break + bytes_read += len(data) + return bytes_read + + +def prefetch_checkpoint_files( + hf_weights_files: list[str], + num_threads: int = 4, +) -> threading.Thread: + """Prefetch checkpoint shards into the OS page cache in the background.""" + sorted_files = sorted(hf_weights_files) + rank, world_size = _get_checkpoint_prefetch_rank_info() + my_files = sorted_files[rank::world_size] + num_threads = max(1, num_threads) + + logger.info( + "Rank %d: prefetching %d/%d checkpoint shards into OS page cache " + "(background, %d local ranks sharing work, %d threads per rank).", + rank, + len(my_files), + len(sorted_files), + world_size, + num_threads, + ) + + def _run_prefetch() -> None: + start = time.perf_counter() + bytes_read = 0 + with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as executor: + futures = [ + executor.submit(_prefetch_checkpoint_file, path) for path in my_files + ] + for future in concurrent.futures.as_completed(futures): + try: + bytes_read += future.result() + except Exception: + logger.warning( + "Failed to prefetch a checkpoint shard.", exc_info=True + ) + + elapsed = time.perf_counter() - start + logger.info( + "Rank %d: checkpoint prefetch finished in %.2fs, %.2f GiB read.", + rank, + elapsed, + bytes_read / (1024**3), + ) + + thread = threading.Thread(target=_run_prefetch, daemon=True) + thread.start() + return thread + + +def safetensors_weights_iterator( + hf_weights_files: list[str], + is_all_weights_sharded: bool = False, + decryption_key: str | None = None, + prefetch: bool = False, + prefetch_num_threads: int = 4, +) -> Generator[tuple[str, torch.Tensor], None, None]: + """Iterate over the weights in the model safetensor files. + + If is_all_weights_sharded is True, it uses more optimize read by reading an + entire file instead of reading each tensor one by one. + """ + if decryption_key: + yield from safetensors_encrypted_weights_iterator( + hf_weights_files, is_all_weights_sharded, decryption_key + ) + return + + enable_tqdm = ( + not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 + ) + if prefetch: + prefetch_checkpoint_files( + hf_weights_files, + num_threads=prefetch_num_threads, + ) + + for st_file in tqdm( + hf_weights_files, + desc="Loading safetensors checkpoint shards", + disable=not enable_tqdm, + bar_format=_BAR_FORMAT, + ): + result = safetensors.torch.load_file(st_file, device="cpu") + yield from result.items() + + +def pt_weights_iterator( + hf_weights_files: list[str], +) -> Generator[tuple[str, torch.Tensor], None, None]: + """Iterate over the weights in the model bin/pt files.""" + enable_tqdm = ( + not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 + ) + for bin_file in tqdm( + hf_weights_files, + desc="Loading pt checkpoint shards", + disable=not enable_tqdm, + bar_format=_BAR_FORMAT, + ): + state = torch.load(bin_file, map_location="cpu") + yield from state.items() + del state + torch.cuda.empty_cache() + + +def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + """Default weight loader.""" + if param.numel() == 1 and loaded_weight.numel() == 1: + # Sometimes scalar values aren't considered tensors with shapes + # so if both param and loaded_weight are a scalar, + # "broadcast" instead of copy + param.data.fill_(loaded_weight.item()) + else: + if param.size() != loaded_weight.size(): + raise ValueError( + f"Attempted to load weight ({loaded_weight.size()}) " + f"into parameter ({param.size()})" + ) + + param.data.copy_(loaded_weight) + + +LoaderFunction = Callable[[torch.Tensor, torch.Tensor], torch.Tensor] + + +def sharded_weight_loader(shard_axis: int, tp_rank: int) -> LoaderFunction: + """Create a weight loader that shards the weights along the given axis""" + + def loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: + shard_size = param.data.shape[shard_axis] + start_idx = tp_rank * shard_size + loaded_weight = loaded_weight.narrow(shard_axis, start_idx, shard_size) + + return default_weight_loader(param, loaded_weight) + + return loader + + +def initialize_dummy_weights( + model: torch.nn.Module, + low: float = -1e-3, + high: float = 1e-3, + seed: int = 1234, +) -> None: + """Initialize model weights with random values. + + The model weights must be randomly initialized for accurate performance + measurements. Additionally, the model weights should not cause NaNs in the + forward pass. We empirically found that initializing the weights with + values between -1e-3 and 1e-3 works well for most models. + + We use per-parameter random seed, so that dummy weights are consistent, + even if the model is partitioned across multiple devices. When the seed + is fixed, the random values generated by this function only depends on + the parameter's number of elements and its data type. + """ + for param in model.state_dict().values(): + if torch.is_floating_point(param): + generator = torch.Generator(device=param.data.device) + generator.manual_seed(seed) + if torch.finfo(param.data.dtype).bits < 16: + # uniform_ doesn't support < 16-bit datatypes (FP8) + dtype = param.data.dtype + tmp_param = param.data.to(torch.float16) + tmp_param = tmp_param.uniform_(low, high, generator=generator).to(dtype) + param.data.copy_(tmp_param) + else: + param.uniform_(low, high, generator=generator) + + +class KVCacheQuantSchema(BaseModel): + dtype: str + # Each key is a TP rank. Each value is a dictionary mapping a TP rank's + # layer indices to their per-tensor KV cache scaling factor. + # own schema class (tricky as its members are variable) + scaling_factor: dict[int, dict[int, float]] + + @model_validator(mode="after") + def check_is_fp8(self) -> "KVCacheQuantSchema": + if self.dtype != "float8_e4m3fn": + raise ValueError( + "Loaded scaling factors intended for KV cache dtype = " + f"{self.dtype} rather than float8_e4m3fn!" + ) + return self + + @model_validator(mode="after") + def check_tp_ranks(self, info: ValidationInfo) -> "KVCacheQuantSchema": + context = info.context + if context: + tp_size = context["tp_size"] + num_hidden_layers = context["num_hidden_layers"] + if len(self.scaling_factor) != tp_size: + raise ValueError( + f"Loaded dictionary has TP size {len(self.scaling_factor)} " + f"but LLM engine is currently running with TP size {tp_size}." + ) + for tp_rank, layer_maps in self.scaling_factor.items(): + if len(layer_maps) != num_hidden_layers: + raise ValueError( + f"KV cache scales map for TP rank {tp_rank} is malformed. " + f"Expected {num_hidden_layers} layers, got " + f"{len(layer_maps)}." + ) + for i in range(tp_size): + if i not in self.scaling_factor: + raise ValueError(f"KV cache scales map for TP rank {i} not found.") + return self + + @model_validator(mode="after") + def check_current_rank(self, info: ValidationInfo) -> "KVCacheQuantSchema": + context = info.context + if context: + tp_rank = context["tp_rank"] + num_hidden_layers = context["num_hidden_layers"] + layer_scales_map = self.scaling_factor[tp_rank] + for i in range(num_hidden_layers): + if i not in layer_scales_map: + raise ValueError( + f"Could not find KV cache scales for layer {i} in " + f"TP rank {tp_rank}." + ) + return self + + +class QuantParamSchema(BaseModel): + # (e.g. weights/activations params) once functionality is enabled + model_config = ConfigDict(protected_namespaces=()) + model_type: str | None + kv_cache: KVCacheQuantSchema + + @model_validator(mode="after") + def check_model_type(self, info: ValidationInfo) -> "QuantParamSchema": + context = info.context + if context: + model_type = context.get("model_type", None) + if model_type is not None: + if model_type != self.model_type: + raise ValueError( + f"Model type is {model_type} but loaded " + f"scaling factors belonging to different " + f"model type {self.model_type}!" + ) + return self + + +def kv_cache_scales_loader( + filename: str, + tp_rank: int, + tp_size: int, + num_hidden_layers: int, + model_type: str | None, +) -> Iterable[tuple[int, float]]: + """ + A simple utility to read in KV cache scaling factors that have been + previously serialized to disk. Used by the model to populate the appropriate + KV cache scaling factors. The serialization should represent a dictionary + whose keys are the TP ranks and values are another dictionary mapping layers + to their KV cache scaling factors. + """ + try: + with open(filename) as f: + context = { + "model_type": model_type, + "num_hidden_layers": num_hidden_layers, + "tp_rank": tp_rank, + "tp_size": tp_size, + } + schema_dct = json.load(f) + schema = QuantParamSchema.model_validate(schema_dct, context=context) + layer_scales_map = schema.kv_cache.scaling_factor[tp_rank] + return layer_scales_map.items() + except FileNotFoundError: + logger.error("File or directory '%s' not found.", filename) + except json.JSONDecodeError: + logger.error("Error decoding JSON in file '%s'.", filename) + except Exception: + logger.error("An error occurred while reading '%s'.", filename) + # This section is reached if and only if any of the excepts are hit + # Return an empty iterable (list) => no KV cache scales are loaded + # which ultimately defaults to 1.0 scales + logger.warning( + "Defaulting to KV cache scaling factors = 1.0 for all " + "layers in TP rank %d as an error occurred during loading.", + tp_rank, + ) + return [] + + +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. + # 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/tokenspeed/runtime/models/base/__init__.py b/python/tokenspeed/runtime/models/base/__init__.py new file mode 100644 index 0000000..a42157e --- /dev/null +++ b/python/tokenspeed/runtime/models/base/__init__.py @@ -0,0 +1,37 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from tokenspeed.runtime.models.base.causal_lm import BaseCausalLM +from tokenspeed.runtime.models.base.decoder_layer import ( + BaseDecoderLayer, + BaseMoEDecoderLayer, + CompiledDecoderLayer, + CompiledMoEDecoderLayer, +) +from tokenspeed.runtime.models.base.transformer_model import BaseTransformerModel + +__all__ = [ + "BaseCausalLM", + "BaseDecoderLayer", + "BaseMoEDecoderLayer", + "CompiledDecoderLayer", + "CompiledMoEDecoderLayer", + "BaseTransformerModel", +] diff --git a/python/tokenspeed/runtime/models/base/causal_lm.py b/python/tokenspeed/runtime/models/base/causal_lm.py new file mode 100644 index 0000000..b0352f6 --- /dev/null +++ b/python/tokenspeed/runtime/models/base/causal_lm.py @@ -0,0 +1,251 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Base causal language model: model + lm_head + logits_processor.""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +import torch +from torch import nn +from transformers import PretrainedConfig + +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.layers.linear import ReplicatedLinear +from tokenspeed.runtime.layers.logits_processor import LogitsMetadata, LogitsProcessor +from tokenspeed.runtime.layers.quantization import QuantizationConfig +from tokenspeed.runtime.layers.vocab_parallel_embedding import ParallelLMHead +from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader +from tokenspeed.runtime.models.base.transformer_model import BaseTransformerModel +from tokenspeed.runtime.utils import add_prefix + + +class BaseCausalLM(nn.Module): + + model_cls: type[BaseTransformerModel] + + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + encoder_only: bool = False, + ) -> None: + + super().__init__() + self.config = config + self.mapping = mapping + self.quant_config = quant_config + self.capture_aux_hidden_states: bool = False + + self.encoder_only = encoder_only + if encoder_only: + # Vision-only role (EPD encode): never allocate the LM / lm_head / + # logits processor (the LM allocation is the OOM at encode TP=1). + # self.config is already set above for the vision path + # (separate_deepstack_embeds needs self.config.hidden_size). + self.model = None + self.lm_head = None + self.logits_processor = None + else: + self.model = self.resolve_model(config, mapping, quant_config, prefix) + self.lm_head = self.resolve_lm_head(config, quant_config, prefix) + self.logits_processor = self.resolve_logits_processor(config) + self.post_init() + + def resolve_model( + self, + config: PretrainedConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None, + prefix: str, + ) -> BaseTransformerModel: + + return self.model_cls( + config, + mapping=mapping, + quant_config=quant_config, + prefix=add_prefix("model", prefix), + ) + + def resolve_lm_head( + self, + config: PretrainedConfig, + quant_config: QuantizationConfig | None, + prefix: str, + ) -> nn.Module: + + if getattr(config, "tie_word_embeddings", False): + return self.model.embed_tokens + + if self.mapping.attn.has_dp: + return ReplicatedLinear( + config.hidden_size, + config.vocab_size, + bias=False, + prefix=add_prefix("lm_head", prefix), + ) + + return ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + + def resolve_logits_processor(self, config: PretrainedConfig) -> LogitsProcessor: + + return LogitsProcessor( + config, + skip_all_gather=self.mapping.attn.has_dp, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + + def post_init(self) -> None: + """Hook for subclasses that need derived state after shared modules exist.""" + + def set_eagle3_layers_to_capture(self, layer_ids: list[int] | None = None) -> None: + + self.capture_aux_hidden_states = True + + if layer_ids is None: + + num_layers = self.config.num_hidden_layers + self.model.layers_to_capture = [2, num_layers // 2, num_layers - 3] + + else: + + self.model.layers_to_capture = [val + 1 for val in layer_ids] + + @torch.no_grad() + def forward( + self, + ctx: ForwardContext, + input_ids: torch.Tensor, + positions: torch.Tensor, + out_cache_loc: torch.Tensor, + **kwargs, + ) -> torch.Tensor: + + model_kwargs = self.prepare_model_kwargs(ctx, input_ids, kwargs) + + hidden_states, aux_hidden_states = self.model( + input_ids, + positions, + ctx, + out_cache_loc, + **model_kwargs, + ) + logits_metadata = LogitsMetadata.from_forward_context(ctx) + + return self.logits_processor( + input_ids, + hidden_states, + self.lm_head, + logits_metadata, + aux_hidden_states, + ) + + def prepare_model_kwargs( + self, ctx: ForwardContext, input_ids: torch.Tensor, kwargs: dict + ) -> dict: + """Hook for subclasses to pass model-specific tensors.""" + model_kwargs = {} + for key in ("input_embeds", "inputs_embeds"): + if kwargs.get(key) is not None: + model_kwargs[key] = kwargs[key] + return model_kwargs + + # Weight loading. + + def get_stacked_params_mapping(self) -> list[tuple[str, str, str]]: + + return [] + + def get_skip_weight_names(self) -> list[str]: + + return ["rotary_emb.inv_freq"] + + def load_weights( + self, weights: Iterable[tuple[str, torch.Tensor]], **kwargs: Any + ) -> None: + + stacked_params_mapping = self.get_stacked_params_mapping() + skip_patterns = self.get_skip_weight_names() + params_dict: dict[str, nn.Parameter] = dict(self.named_parameters()) + + for name, loaded_weight in weights: + + if any(pattern in name for pattern in skip_patterns): + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + + if weight_name not in name: + continue + + name = name.replace(weight_name, param_name) + + if name.endswith(".bias") and name not in params_dict: + continue + + if name not in params_dict: + continue + + param = params_dict[name] + param.weight_loader(param, loaded_weight, shard_id) + + break + + else: + + if name.endswith(".bias") and name not in params_dict: + continue + + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + def get_embed_and_head(self) -> tuple[torch.Tensor, torch.Tensor]: + + return self.model.embed_tokens.weight, self.lm_head.weight + + def set_embed_and_head(self, embed: torch.Tensor, head: torch.Tensor) -> None: + + del self.model.embed_tokens.weight + del self.lm_head.weight + + self.model.embed_tokens.weight = embed + self.lm_head.weight = head + + torch.cuda.empty_cache() + torch.cuda.synchronize() diff --git a/python/tokenspeed/runtime/models/base/comm_ops.py b/python/tokenspeed/runtime/models/base/comm_ops.py new file mode 100644 index 0000000..db5e722 --- /dev/null +++ b/python/tokenspeed/runtime/models/base/comm_ops.py @@ -0,0 +1,412 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""CommOp: communication operations automatically inserted by the layer compiler. + +Each ``CommOp`` is an ``nn.Module`` that performs a single communication +primitive (all-reduce, reduce-scatter, all-gather, or fused variants). +They are created by the compiler based on Placement transitions between +adjacent compute modules. +""" + +from __future__ import annotations + +import torch +from torch import nn + +from tokenspeed.runtime.distributed.comm_ops import ( + all_reduce, + token_all_gather, + token_reduce_scatter, +) +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.models.base.placement import ParallelGroup + +# --------------------------------------------------------------------------- +# Helpers for computing scattered token counts from ForwardContext +# --------------------------------------------------------------------------- + + +def _scatter_count(num_tokens: int, tp_size: int) -> list[int]: + base, remainder = divmod(num_tokens, tp_size) + return [base + 1] * remainder + [base] * (tp_size - remainder) + + +def _scattered_num_tokens_all(ctx: ForwardContext, mapping: Mapping) -> list[int]: + if ctx.global_num_tokens is not None: + scattered: list[int] = [] + for attn_dp_rank in range(mapping.attn.dp_size): + # global_num_tokens is indexed by global rank with dp stride + # tp_size * cp_size; cp peers report the same count. + num_tokens = ctx.global_num_tokens[ + attn_dp_rank * mapping.attn.tp_size * mapping.attn.cp_size + ] + scattered.extend(_scatter_count(num_tokens, mapping.attn.tp_size)) + return scattered + return _scatter_count(ctx.input_num_tokens, mapping.attn.tp_size) + + +def _group_scattered_num_tokens( + ctx: ForwardContext, + mapping: Mapping, + group_type: ParallelGroup, +) -> list[int]: + if group_type == ParallelGroup.ATTN_TP: + start = mapping.attn.tp_size * mapping.attn.dp_rank + end = start + mapping.attn.tp_size + return _scattered_num_tokens_all(ctx, mapping)[start:end] + elif group_type == ParallelGroup.DENSE_TP: + start = mapping.dense.tp_size * mapping.dense.dp_rank + end = start + mapping.dense.tp_size + return _scattered_num_tokens_all(ctx, mapping)[start:end] + elif group_type == ParallelGroup.MOE_TP_EP: + tp_ep_size = mapping.moe.tp_ep_size + # Without DP, all ranks share the batch and the scattered table needs + # no global metadata, so the lookup below stays valid. + if ctx.global_num_tokens is not None or not mapping.attn.has_dp: + # After the attention reduce-scatter, each rank holds its + # scattered share of its attn dp group's tokens, not the raw + # global count; MoE collectives must size from those rows. + scattered = _scattered_num_tokens_all(ctx, mapping) + return [ + scattered[mapping.attn.scatter_index(rank)] + for rank in mapping.moe.tp_ep_group + ] + # With DP but no gathered metadata, other dp groups' counts are + # unknown; only the local rank's contribution can be reported. + result = [0] * tp_ep_size + result[mapping.moe.tp_ep_rank] = ctx.input_num_tokens + return result + else: + raise ValueError(f"Unknown parallel group type: {group_type}") + + +# --------------------------------------------------------------------------- +# Group info +# --------------------------------------------------------------------------- + + +def _get_group_info( + mapping: Mapping, group_type: ParallelGroup +) -> tuple[int, tuple[int, ...], bool]: + """Return (rank, group, has_parallelism) for the given parallel group type.""" + if group_type == ParallelGroup.ATTN_TP: + return mapping.attn.tp_rank, mapping.attn.tp_group, mapping.has_attn_tp + elif group_type == ParallelGroup.DENSE_TP: + return mapping.dense.tp_rank, mapping.dense.tp_group, mapping.dense.has_tp + elif group_type == ParallelGroup.MOE_TP_EP: + return mapping.moe.tp_ep_rank, mapping.moe.tp_ep_group, mapping.moe.has_tp_ep + else: + raise ValueError(f"Unknown parallel group type: {group_type}") + + +def _should_fuse_allreduce_norm( + num_tokens: int, + *, + has_parallel: bool, + use_all_reduce_mode: bool = True, +) -> bool: + from tokenspeed.runtime.utils.env import global_server_args_dict + + return ( + use_all_reduce_mode + and has_parallel + and global_server_args_dict.get("enable_allreduce_fusion", False) + and num_tokens > 0 + and num_tokens <= global_server_args_dict["comm_fusion_max_num_tokens"] + ) + + +# --------------------------------------------------------------------------- +# Communication Operations +# --------------------------------------------------------------------------- + + +class CommOp(nn.Module): + """Base class for compiler-inserted communication operations.""" + + def __init__(self, mapping: Mapping, group_type: ParallelGroup) -> None: + super().__init__() + self.mapping = mapping + self.group_type = group_type + rank, group, has_parallel = _get_group_info(mapping, group_type) + self._rank = rank + self._group = group + self._has_parallel = has_parallel + + +class AllReduceOp(CommOp): + """all_reduce: Partial -> Replicate.""" + + def forward( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ctx: ForwardContext, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if not self._has_parallel: + return hidden_states, residual + hidden_states = all_reduce(hidden_states, self._group) + return hidden_states, residual + + +class ReduceScatterOp(CommOp): + """reduce_scatter: Partial -> Shard.""" + + def forward( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ctx: ForwardContext, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if not self._has_parallel: + return hidden_states, residual + scattered_num_tokens = _group_scattered_num_tokens( + ctx, self.mapping, self.group_type + ) + hidden_states = token_reduce_scatter( + hidden_states, + group=self._group, + scattered_num_tokens=scattered_num_tokens, + ) + return hidden_states, residual + + +class AllGatherOp(CommOp): + """all_gather: Shard -> Replicate.""" + + def forward( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ctx: ForwardContext, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if not self._has_parallel: + return hidden_states, residual + scattered_num_tokens = _group_scattered_num_tokens( + ctx, self.mapping, self.group_type + ) + hidden_states = token_all_gather( + hidden_states, + group=self._group, + scattered_num_tokens=scattered_num_tokens, + ) + return hidden_states, residual + + +class ResidualAllGatherOp(CommOp): + """all_gather the residual: needed when transitioning from RSAG -> AR mode.""" + + def forward( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ctx: ForwardContext, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if not self._has_parallel or residual is None: + return hidden_states, residual + scattered_num_tokens = _group_scattered_num_tokens( + ctx, self.mapping, self.group_type + ) + residual = token_all_gather( + residual, + group=self._group, + scattered_num_tokens=scattered_num_tokens, + ) + return hidden_states, residual + + +class ResidualSliceOp(CommOp): + """Slice residual when transitioning from AR -> RSAG mode. + + When the previous layer used all-reduce (residual has full tokens) but the + current layer uses reduce-scatter (residual should be scattered), we need + to slice the residual to keep only the local portion. + """ + + def forward( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ctx: ForwardContext, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if not self._has_parallel or residual is None: + return hidden_states, residual + scattered_num_tokens = _group_scattered_num_tokens( + ctx, self.mapping, self.group_type + ) + offset = sum(scattered_num_tokens[: self._rank]) + residual = residual[offset : offset + hidden_states.size(0)] + return hidden_states, residual + + +class FusedReduceNormOp(CommOp): + """Fused allreduce + residual + RMSNorm. + + When conditions are met (all-reduce mode, small enough token count), this + replaces separate allreduce + norm with a single fused kernel. Falls back + to unfused path when fusion is not beneficial. + """ + + def __init__( + self, + mapping: Mapping, + group_type: ParallelGroup, + norm_module: nn.Module, + ) -> None: + super().__init__(mapping, group_type) + self.norm_module = norm_module + + def _should_fuse(self, num_tokens: int) -> bool: + return _should_fuse_allreduce_norm( + num_tokens, + has_parallel=self._has_parallel, + ) + + def forward( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ctx: ForwardContext, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if residual is None: + # First layer: no residual to fuse with, just norm + residual = hidden_states + hidden_states = self.norm_module(hidden_states) + return hidden_states, residual + + if self._should_fuse(hidden_states.shape[0]): + hidden_states, residual, *_ = ( + self.norm_module.forward_with_allreduce_fusion( + self._rank, + self._group, + hidden_states, + residual, + ) + ) + else: + # Fusion not available — fall back to explicit allreduce + norm. + # The hidden_states arriving here are Partial (unreduced) from + # the preceding compute module's output. We must allreduce + # before applying the norm. + if self._has_parallel: + hidden_states = all_reduce(hidden_states, self._group) + hidden_states, residual = self.norm_module(hidden_states, residual) + return hidden_states, residual + + +class DeferredReduceOp(CommOp): + """A marker that indicates allreduce is deferred to the downstream norm op. + + The reduce is always deferred — the downstream ``FusedReduceNormOp`` or + ``FinalNormOp`` is responsible for performing the all-reduce (fused or + explicit) before applying the norm. This op is therefore a no-op at + runtime; it exists so that the compiler can record the deferred state. + """ + + def forward( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ctx: ForwardContext, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + # Always defer — the downstream norm op handles the reduce. + return hidden_states, residual + + +class FinalNormOp(CommOp): + """Final norm after last layer, optionally fusing deferred allreduce. + + Also handles the post-final-norm all-gather needed in RSAG mode for the + LM head. + """ + + def __init__( + self, + mapping: Mapping, + group_type: ParallelGroup, + norm_module: nn.Module, + use_all_reduce_mode: bool, + lm_head_group_type: ParallelGroup | None = None, + ) -> None: + super().__init__(mapping, group_type) + self.norm_module = norm_module + self.use_all_reduce_mode = use_all_reduce_mode + # The LM head follows attn_tp sharding, so in RSAG mode the + # all-gather must use the attn_tp group — which may differ from + # group_type (e.g. when the last layer outputs on DENSE_TP). + if lm_head_group_type is not None and lm_head_group_type != group_type: + lm_rank, lm_group, lm_has_parallel = _get_group_info( + mapping, lm_head_group_type + ) + self._lm_head_group_type = lm_head_group_type + self._lm_rank = lm_rank + self._lm_group = lm_group + self._lm_has_parallel = lm_has_parallel + else: + self._lm_head_group_type = group_type + self._lm_rank = self._rank + self._lm_group = self._group + self._lm_has_parallel = self._has_parallel + + def _should_fuse(self, num_tokens: int) -> bool: + return _should_fuse_allreduce_norm( + num_tokens, + has_parallel=self._has_parallel, + use_all_reduce_mode=self.use_all_reduce_mode, + ) + + def forward( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ctx: ForwardContext, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + # Returns (normed hidden states, post-add residual); see + # CommManager.final_norm for the residual's meaning. + if self._should_fuse(hidden_states.shape[0]): + hidden_states, residual_out, *_ = ( + self.norm_module.forward_with_allreduce_fusion( + self._rank, + self._group, + hidden_states, + residual, + ) + ) + else: + # The preceding DeferredReduceOp always defers, so we must + # perform the all-reduce here before applying the norm. + if self._has_parallel and self.use_all_reduce_mode: + hidden_states = all_reduce(hidden_states, self._group) + hidden_states, residual_out = self.norm_module(hidden_states, residual) + # In RSAG mode, all-gather to restore tokens for the LM head. + # Uses the LM head group (ATTN_TP) which may differ from the + # scatter group when attn_tp != dense_tp. + if self._lm_has_parallel and not self.use_all_reduce_mode: + scattered_num_tokens = _group_scattered_num_tokens( + ctx, self.mapping, self._lm_head_group_type + ) + hidden_states = token_all_gather( + hidden_states, + group=self._lm_group, + scattered_num_tokens=scattered_num_tokens, + ) + return hidden_states, residual_out diff --git a/python/tokenspeed/runtime/models/base/compiler.py b/python/tokenspeed/runtime/models/base/compiler.py new file mode 100644 index 0000000..e143df8 --- /dev/null +++ b/python/tokenspeed/runtime/models/base/compiler.py @@ -0,0 +1,557 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Layer Compiler: analyses ModuleSpec annotations and inserts CommOps. + +The compiler inspects each decoder layer's sub-modules (in the order declared +by ``resolve_exec_plan``), examines adjacent Placement pairs, and inserts the +minimal set of communication operations to transition between them. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +from torch import nn + +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.models.base.comm_ops import ( + AllGatherOp, + AllReduceOp, + CommOp, + DeferredReduceOp, + FusedReduceNormOp, + ReduceScatterOp, + ResidualAllGatherOp, + ResidualSliceOp, + _scattered_num_tokens_all, +) +from tokenspeed.runtime.models.base.execution import ( + CompiledDecoderLayer, + ExecutionNode, + ExecutionState, + ExecutionStep, + StepRunner, +) +from tokenspeed.runtime.models.base.module_spec import ( + CallConvention, + FusionCapability, + ModuleKind, + ModuleSpec, +) +from tokenspeed.runtime.models.base.placement import ( + ParallelGroup, + Partial, + Placement, + PlacementType, + Replicate, + Shard, + can_fuse_reduce_norm, + group_has_parallel, + use_all_reduce, +) + + +@dataclass +class _TrackedState: + hidden: Placement | None = None + residual: Placement | None = None + + +# --------------------------------------------------------------------------- +# Runner factories — dispatch on CallConvention +# --------------------------------------------------------------------------- + + +def _runner_from_node(node: ExecutionNode) -> StepRunner: + module = node.module + spec = node.spec + + if isinstance(module, FusedReduceNormOp): + return lambda state, positions: _run_fused_reduce_norm(module, state) + if spec.call == CallConvention.NORM_WITH_OPTIONAL_RESIDUAL: + return lambda state, positions: _run_norm(module, state) + if spec.call == CallConvention.ATTENTION: + return lambda state, positions: _run_attention(module, state, positions) + if spec.call == CallConvention.MOE: + return lambda state, positions: _run_moe(module, state) + return lambda state, positions: _run_hidden_states_only(module, state) + + +# --------------------------------------------------------------------------- +# Per-convention runner functions +# --------------------------------------------------------------------------- + + +def _run_fused_reduce_norm( + module: FusedReduceNormOp, + state: ExecutionState, +) -> ExecutionState: + hidden_states, residual = module(state.hidden_states, state.residual, state.ctx) + return ExecutionState(hidden_states, residual, state.ctx, state.out_cache_loc) + + +def _run_norm(module: nn.Module, state: ExecutionState) -> ExecutionState: + if state.residual is not None: + hidden_states, residual = module(state.hidden_states, state.residual) + else: + residual = state.hidden_states + hidden_states = module(state.hidden_states) + return ExecutionState(hidden_states, residual, state.ctx, state.out_cache_loc) + + +def _run_attention( + module: nn.Module, + state: ExecutionState, + positions: torch.Tensor, +) -> ExecutionState: + hidden_states = module( + positions=positions, + hidden_states=state.hidden_states, + ctx=state.ctx, + out_cache_loc=state.out_cache_loc, + ) + return ExecutionState(hidden_states, state.residual, state.ctx, state.out_cache_loc) + + +def _run_moe(module: nn.Module, state: ExecutionState) -> ExecutionState: + scattered = _scattered_num_tokens_all(state.ctx, module.mapping) + num_global_tokens = sum(scattered) + max_num_tokens_per_gpu = max(scattered) if scattered else 0 + hidden_states = module( + state.hidden_states, num_global_tokens, max_num_tokens_per_gpu + ) + return ExecutionState(hidden_states, state.residual, state.ctx, state.out_cache_loc) + + +def _run_hidden_states_only(module: nn.Module, state: ExecutionState) -> ExecutionState: + hidden_states = module(state.hidden_states) + return ExecutionState(hidden_states, state.residual, state.ctx, state.out_cache_loc) + + +# --------------------------------------------------------------------------- +# Placement helpers +# --------------------------------------------------------------------------- + + +def _input_group(spec: ModuleSpec) -> ParallelGroup | None: + return spec.input_placement.group if spec.input_placement else None + + +def _output_group(spec: ModuleSpec) -> ParallelGroup | None: + return spec.output_placement.group if spec.output_placement else None + + +def _find_last_compute_index(exec_plan: list[ExecutionNode]) -> int: + """Find the index of the last compute (non-NORM) module in exec_plan.""" + last_idx = -1 + for i, mod in enumerate(exec_plan): + spec = mod.spec + if spec.kind != ModuleKind.NORM: + last_idx = i + return last_idx + + +# --------------------------------------------------------------------------- +# Main compilation entry point +# --------------------------------------------------------------------------- + + +def compile_decoder_layer( + layer: nn.Module, + exec_plan: list[ExecutionNode], + mapping: Mapping, + prev_layer_output_group: ParallelGroup | None = None, + next_layer_input_group: ParallelGroup | None = None, +) -> CompiledDecoderLayer: + """Analyse a decoder layer execution plan and produce a CompiledDecoderLayer.""" + + last_compute_idx = _find_last_compute_index(exec_plan) + + steps: list[ExecutionStep] = [] + + first_compute_input_group = find_first_compute_input_group(exec_plan) + state = _TrackedState( + hidden=_initial_hidden_placement( + mapping=mapping, + prev_layer_output_group=prev_layer_output_group, + first_compute_input_group=first_compute_input_group, + ), + residual=_initial_residual_placement( + mapping=mapping, + prev_layer_output_group=prev_layer_output_group, + first_compute_input_group=first_compute_input_group, + ), + ) + + is_first_layer = prev_layer_output_group is None + + for mod_idx, node in enumerate(exec_plan): + spec = node.spec + is_last_compute = mod_idx == last_compute_idx + + if spec.kind == ModuleKind.NORM: + step = _compile_norm_step( + node=node, + mapping=mapping, + mod_idx=mod_idx, + exec_plan=exec_plan, + state=state, + ) + steps.append(step) + else: + step = _compile_compute_step( + node=node, + mapping=mapping, + mod_idx=mod_idx, + exec_plan=exec_plan, + is_last_compute=is_last_compute, + state=state, + next_layer_input_group=next_layer_input_group, + is_first_layer=is_first_layer, + ) + steps.append(step) + + # Determine the final placement after this layer. + final_placement = _compute_final_placement(state, mapping) + + return CompiledDecoderLayer( + steps=steps, + final_placement=final_placement, + mapping=mapping, + ) + + +# --------------------------------------------------------------------------- +# Per-step compilation +# --------------------------------------------------------------------------- + + +def _compile_norm_step( + node: ExecutionNode, + mapping: Mapping, + mod_idx: int, + exec_plan: list[ExecutionNode], + state: _TrackedState, +) -> ExecutionStep: + """Compile a NORM step.""" + module = node.module + spec = node.spec + next_compute_group = _find_next_compute_input_group(exec_plan, mod_idx) + hidden = state.hidden + src_group = hidden.group if hidden is not None else None + prev_output_is_partial = hidden is not None and hidden.type == PlacementType.PARTIAL + + if ( + prev_output_is_partial + and spec.fusion == FusionCapability.REDUCE_NORM + and src_group is not None + and next_compute_group is not None + and can_fuse_reduce_norm(mapping, src_group, next_compute_group) + ): + fused_norm = FusedReduceNormOp(mapping, src_group, module) + state.hidden = Replicate(src_group) + fused_node = ExecutionNode( + module=fused_norm, + spec=spec, + name=node.name, + ) + return ExecutionStep( + runner=_runner_from_node(fused_node), + module=fused_norm, + spec=spec, + kind=spec.kind, + captures_aux=spec.captures_aux, + skip_on_idle=spec.skip_on_idle, + name=node.name, + ) + else: + return ExecutionStep( + runner=_runner_from_node(node), + module=module, + spec=spec, + kind=spec.kind, + captures_aux=spec.captures_aux, + skip_on_idle=spec.skip_on_idle, + name=node.name, + ) + + +def _compile_compute_step( + node: ExecutionNode, + mapping: Mapping, + mod_idx: int, + exec_plan: list[ExecutionNode], + is_last_compute: bool, + state: _TrackedState, + next_layer_input_group: ParallelGroup | None, + is_first_layer: bool = False, +) -> ExecutionStep: + """Compile a compute step (ATTENTION / DENSE_MLP / MOE / GENERIC).""" + spec = node.spec + pre_comms: list[CommOp] = [] + post_comms: list[CommOp] = [] + + input_group = _input_group(spec) + output_group = _output_group(spec) + + hidden = state.hidden + + if ( + spec.input_placement is not None + and spec.input_placement.type == PlacementType.REPLICATE + and input_group is not None + and group_has_parallel(mapping, input_group) + and hidden is not None + and hidden.type == PlacementType.SHARD + ): + gather_group = hidden.group + pre_comms.append(AllGatherOp(mapping, gather_group)) + if state.residual is not None and state.residual.type == PlacementType.SHARD: + pre_comms.append(ResidualAllGatherOp(mapping, gather_group)) + state.residual = Replicate(input_group) + state.hidden = Replicate(input_group) + elif hidden is None and not (is_first_layer and spec.kind == ModuleKind.ATTENTION): + # Data is not tracked (no previous TP/EP), but the current module + # expects Replicate on a group with compiler-managed parallelism + # (e.g. Dense TP, MoE TP/EP). All-gather on the input group. + # The first layer's attention is exempt: data from embedding. + pre_comms.append(AllGatherOp(mapping, input_group)) + state.hidden = Replicate(input_group) + + residual_before_post = state.residual + + state.hidden = ( + Partial(output_group) + if output_group is not None and group_has_parallel(mapping, output_group) + else None + ) + + if is_last_compute: + _insert_last_compute_post_comms( + post_comms=post_comms, + spec=spec, + mapping=mapping, + next_layer_input_group=next_layer_input_group, + exec_plan=exec_plan, + state=state, + hidden_before_input=hidden, + residual_before_output=residual_before_post, + ) + else: + _insert_mid_layer_post_comms( + post_comms=post_comms, + spec=spec, + mapping=mapping, + mod_idx=mod_idx, + exec_plan=exec_plan, + state=state, + ) + + return ExecutionStep( + runner=_runner_from_node(node), + pre_comms=pre_comms, + post_comms=post_comms, + spec=spec, + kind=spec.kind, + captures_aux=spec.captures_aux, + skip_on_idle=spec.skip_on_idle, + name=node.name, + ) + + +# --------------------------------------------------------------------------- +# Post-communication insertion +# --------------------------------------------------------------------------- + + +def _insert_last_compute_post_comms( + post_comms: list[CommOp], + spec: ModuleSpec, + mapping: Mapping, + next_layer_input_group: ParallelGroup | None, + exec_plan: list[ExecutionNode], + state: _TrackedState, + hidden_before_input: Placement | None, + residual_before_output: Placement | None, +) -> None: + """Insert post-communication for the last compute module in the layer.""" + output_group = _output_group(spec) + if output_group is None or not group_has_parallel(mapping, output_group): + state.hidden = None + return + + if next_layer_input_group is None: + next_layer_input_group = find_first_compute_input_group(exec_plan) + # AR/RSAG decision must compare against ATTN_TP, not whatever + # group attn_spec may have switched to (e.g. DENSE_TP). + # This matches CommManager.use_all_reduce(is_moe) which checks + # attn.tp_size against the output tp/ep size. + use_ar = use_all_reduce(mapping, output_group, ParallelGroup.ATTN_TP) + first_layer_dense_tp_from_dp_attention = ( + spec.kind == ModuleKind.DENSE_MLP + and hidden_before_input is None + and residual_before_output is None + and not mapping.has_attn_tp + ) + + if not use_ar and first_layer_dense_tp_from_dp_attention: + state.hidden = Shard(output_group) + return + + if use_ar and mapping.has_attn_tp: + post_comms.append(DeferredReduceOp(mapping, output_group)) + state.hidden = Partial(output_group) + elif use_ar: + post_comms.append(AllReduceOp(mapping, output_group)) + state.hidden = Replicate(output_group) + else: + post_comms.append(ReduceScatterOp(mapping, output_group)) + state.hidden = Shard(output_group) + + +def _insert_mid_layer_post_comms( + post_comms: list[CommOp], + spec: ModuleSpec, + mapping: Mapping, + mod_idx: int, + exec_plan: list[ExecutionNode], + state: _TrackedState, +) -> None: + """Insert post-communication for a mid-layer compute module.""" + output_group = _output_group(spec) + if output_group is None or not group_has_parallel(mapping, output_group): + # No TP on this group → output is effectively Replicate, no comm. + state.hidden = Replicate(output_group) if output_group is not None else None + return + + next_compute_input_group = _find_next_compute_input_group(exec_plan, mod_idx) + if next_compute_input_group is None: + return + + next_norm_can_fuse = _intervening_norm_supports_fusion(exec_plan, mod_idx) + use_ar = use_all_reduce(mapping, output_group, next_compute_input_group) + + if next_norm_can_fuse and can_fuse_reduce_norm( + mapping, output_group, next_compute_input_group + ): + # Fused norm will absorb the reduce. + # Data stays Partial until norm resolves it (scattered_on stays None). + return + + if use_ar: + post_comms.append(AllReduceOp(mapping, output_group)) + state.hidden = Replicate(output_group) + if state.residual is not None and state.residual.type == PlacementType.SHARD: + post_comms.append(ResidualAllGatherOp(mapping, output_group)) + state.residual = Replicate(output_group) + return + + post_comms.append(ReduceScatterOp(mapping, output_group)) + state.hidden = Shard(output_group) + if state.residual is None or state.residual.type == PlacementType.REPLICATE: + post_comms.append(ResidualSliceOp(mapping, output_group)) + state.residual = Shard(output_group) + + +# --------------------------------------------------------------------------- +# Placement analysis helpers +# --------------------------------------------------------------------------- + + +def find_first_compute_input_group(exec_plan: list[ExecutionNode]) -> ParallelGroup: + """Find the input group of the first compute (non-NORM) module in exec_plan.""" + for mod in exec_plan: + spec = mod.spec + if spec.kind != ModuleKind.NORM and spec.input_placement is not None: + return spec.input_placement.group + return ParallelGroup.ATTN_TP # fallback + + +def _initial_hidden_placement( + mapping: Mapping, + prev_layer_output_group: ParallelGroup | None, + first_compute_input_group: ParallelGroup, +) -> Placement | None: + if prev_layer_output_group is None: + return None + if not group_has_parallel(mapping, prev_layer_output_group): + return None + # AR/RSAG decision must compare against ATTN_TP, matching + # CommManager.use_all_reduce() which checks attn.tp_size. + if use_all_reduce(mapping, prev_layer_output_group, ParallelGroup.ATTN_TP): + return ( + Partial(prev_layer_output_group) + if mapping.has_attn_tp + else Replicate(prev_layer_output_group) + ) + return Shard(prev_layer_output_group) + + +def _initial_residual_placement( + mapping: Mapping, + prev_layer_output_group: ParallelGroup | None, + first_compute_input_group: ParallelGroup, +) -> Placement | None: + if prev_layer_output_group is None: + return None + if not group_has_parallel(mapping, prev_layer_output_group): + return None + if use_all_reduce(mapping, prev_layer_output_group, ParallelGroup.ATTN_TP): + return Replicate(prev_layer_output_group) + return Shard(prev_layer_output_group) + + +def _find_next_compute_input_group( + exec_plan: list[ExecutionNode], + after_index: int, +) -> ParallelGroup | None: + """Find the input group of the next compute module after *after_index*.""" + for i in range(after_index + 1, len(exec_plan)): + spec = exec_plan[i].spec + if spec.kind != ModuleKind.NORM and spec.input_placement is not None: + return _input_group(spec) + return None + + +def _intervening_norm_supports_fusion( + exec_plan: list[ExecutionNode], + compute_index: int, +) -> bool: + """Check if there's a fusible norm between compute_index and the next compute module.""" + for i in range(compute_index + 1, len(exec_plan)): + spec = exec_plan[i].spec + if spec.kind == ModuleKind.NORM: + return spec.fusion == FusionCapability.REDUCE_NORM + else: + return False + return False + + +def _compute_final_placement( + state: _TrackedState, + mapping: Mapping, +) -> Placement | None: + """Determine the final Placement based on tracked state.""" + hidden = state.hidden + if hidden is None: + return None + return hidden if group_has_parallel(mapping, hidden.group) else hidden diff --git a/python/tokenspeed/runtime/models/base/decoder_layer.py b/python/tokenspeed/runtime/models/base/decoder_layer.py new file mode 100644 index 0000000..ac2c57b --- /dev/null +++ b/python/tokenspeed/runtime/models/base/decoder_layer.py @@ -0,0 +1,369 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Base decoder layer classes. + +``BaseDecoderLayer`` uses CommManager for communication (the default path). +``CompiledDecoderLayer`` uses the compiler-driven path. +""" + +from __future__ import annotations + +from typing import Generic, TypeVar + +import torch +from torch import nn +from transformers import PretrainedConfig + +from tokenspeed.runtime.distributed.comm_manager import CommManager +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.layers.layernorm import RMSNorm +from tokenspeed.runtime.layers.quantization import QuantizationConfig as Q +from tokenspeed.runtime.models.base.execution import ( + CompiledDecoderLayer as _CompiledRuntime, +) +from tokenspeed.runtime.models.base.execution import ( + ExecutionNode, +) +from tokenspeed.runtime.models.base.module_spec import ModuleKind, ModuleSpec +from tokenspeed.runtime.models.base.placement import ParallelGroup, Partial, Replicate + + +def _default_compute_output_placement( + mapping: Mapping, + group: ParallelGroup, +) -> Partial | None: + if group == ParallelGroup.ATTN_TP: + has_parallel = mapping.has_attn_tp + elif group == ParallelGroup.DENSE_TP: + has_parallel = mapping.dense.has_tp + elif group == ParallelGroup.MOE_TP_EP: + has_parallel = mapping.moe.has_tp_ep + else: + raise ValueError(f"Unknown group: {group}") + return Partial(group) if has_parallel else None + + +_C = TypeVar("_C", bound=PretrainedConfig) + + +class BaseDecoderLayer(nn.Module, Generic[_C]): + """Default decoder layer using CommManager for communication. + + Subclasses override ``resolve_attn()`` and ``resolve_mlp()``. + """ + + def __init__( + self, + config: _C, + layer_id: int, + mapping: Mapping, + quant_config: Q | None = None, + prefix: str = "", + ) -> None: + + super().__init__() + + self.config = config + self.quant_config = quant_config + self.layer_id = layer_id + self.total_layers = config.num_hidden_layers + self.mapping = mapping + + self.input_layernorm = self.resolve_norm() + self.post_attention_layernorm = self.resolve_norm() + + self.self_attn = self.resolve_attn(prefix) + self.mlp = self.resolve_mlp(prefix) + + self.comm_manager = CommManager( + mapping=self.mapping, + layer_id=layer_id, + is_moe=self.is_moe_layer, + prev_is_moe=self.is_moe_layer, + input_layernorm=self.input_layernorm, + post_attn_layernorm=self.post_attention_layernorm, + ) + + @property + def is_moe_layer(self) -> bool: + + return False + + def resolve_norm(self) -> nn.Module: + + return RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps) + + def resolve_attn(self, prefix: str) -> nn.Module: + + raise NotImplementedError + + def resolve_mlp(self, prefix: str) -> nn.Module: + + raise NotImplementedError + + def forward_attn( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + residual: torch.Tensor | None, + aux_hidden_states: list | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + + hidden_states, residual = self.comm_manager.input_reduce_norm( + hidden_states, residual + ) + + if aux_hidden_states is not None: + # Under RSAG the residual entering this layer is reduce-scattered + # across the attn TP group; aux consumers (e.g. the EAGLE3 + # drafter) expect full rows, so gather before capturing. + aux_hidden_states.append( + self.comm_manager.gather_residual(residual, ctx).clone() + ) + + hidden_states = self.comm_manager.pre_attn_comm(hidden_states, ctx) + + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ctx=ctx, + out_cache_loc=out_cache_loc, + ) + + hidden_states, residual = self.comm_manager.post_attn_reduce_norm( + hidden_states, residual, ctx + ) + + return hidden_states, residual + + def forward_mlp( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor, + ctx: ForwardContext, + num_global_tokens: int, + max_num_tokens_per_gpu: int, + ) -> torch.Tensor: + + hidden_states = self.comm_manager.pre_mlp_comm(hidden_states, ctx) + + if self.is_moe_layer: + + hidden_states = self.mlp( + hidden_states, num_global_tokens, max_num_tokens_per_gpu + ) + + else: + + hidden_states = self.mlp(hidden_states) + + hidden_states, residual = self.comm_manager.post_mlp_fused( + hidden_states, residual, ctx + ) + + return hidden_states + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + residual: torch.Tensor | None, + aux_hidden_states: list | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + + num_global_tokens, max_num_tokens_per_gpu = self.comm_manager.get_num_tokens( + ctx + ) + + if not ctx.forward_mode.is_idle(): + + hidden_states, residual = self.forward_attn( + positions, + hidden_states, + ctx, + out_cache_loc, + residual, + aux_hidden_states, + ) + + hidden_states = self.forward_mlp( + hidden_states, + residual, + ctx, + num_global_tokens, + max_num_tokens_per_gpu, + ) + + else: + + hidden_states = self.forward_mlp( + hidden_states, + residual, + ctx, + num_global_tokens, + max_num_tokens_per_gpu, + ) + + return hidden_states, residual + + +class BaseMoEDecoderLayer(BaseDecoderLayer): + + @property + def is_moe_layer(self) -> bool: + + return True + + +class CompiledDecoderLayer(nn.Module, Generic[_C]): + """Compiler-driven decoder layer (opt-in). + + Instead of CommManager, the forward delegates to a + ``_CompiledRuntime`` produced by the layer compiler. + """ + + def __init__( + self, + config: _C, + layer_id: int, + mapping: Mapping, + quant_config: Q | None = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.config = config + self.quant_config = quant_config + self.layer_id = layer_id + self.total_layers = config.num_hidden_layers + self.mapping = mapping + self.prefix = prefix + + self._compiled: _CompiledRuntime | None = None + self._exec_plan = self.build_execution_plan(prefix) + + @property + def is_moe_layer(self) -> bool: + return False + + def resolve_norm(self) -> nn.Module: + return RMSNorm(self.config.hidden_size, eps=self.config.rms_norm_eps) + + def build_execution_plan(self, prefix: str) -> list[ExecutionNode]: + self.input_layernorm = self.resolve_norm() + self.self_attn = self.resolve_attn(prefix) + self.post_attention_layernorm = self.resolve_norm() + self.mlp = self.resolve_mlp(prefix) + + return [ + ExecutionNode( + module=self.input_layernorm, + spec=self.norm_spec(captures_aux=True, skip_on_idle=True), + name="input_layernorm", + ), + ExecutionNode( + module=self.self_attn, + spec=self.attn_spec(), + name="self_attn", + ), + ExecutionNode( + module=self.post_attention_layernorm, + spec=self.norm_spec(), + name="post_attention_layernorm", + ), + ExecutionNode( + module=self.mlp, + spec=self.mlp_spec(), + name="mlp", + ), + ] + + def norm_spec( + self, + *, + captures_aux: bool = False, + skip_on_idle: bool = False, + ) -> ModuleSpec: + return ModuleSpec.from_kind( + kind=ModuleKind.NORM, + supports_fused_reduce_norm=True, + captures_aux=captures_aux, + skip_on_idle=skip_on_idle, + ) + + def attn_spec(self) -> ModuleSpec: + input_placement = Replicate(ParallelGroup.ATTN_TP) + return ModuleSpec.from_kind( + input_placement=input_placement, + output_placement=_default_compute_output_placement( + self.mapping, ParallelGroup.ATTN_TP + ), + kind=ModuleKind.ATTENTION, + skip_on_idle=True, + ) + + def mlp_spec(self) -> ModuleSpec: + mlp_group = ( + ParallelGroup.MOE_TP_EP if self.is_moe_layer else ParallelGroup.DENSE_TP + ) + kind = ModuleKind.MOE if self.is_moe_layer else ModuleKind.DENSE_MLP + return ModuleSpec.from_kind( + input_placement=Replicate(mlp_group), + output_placement=_default_compute_output_placement(self.mapping, mlp_group), + kind=kind, + ) + + def resolve_attn(self, prefix: str) -> nn.Module: + raise NotImplementedError + + def resolve_mlp(self, prefix: str) -> nn.Module: + raise NotImplementedError + + def resolve_exec_plan(self) -> list[ExecutionNode]: + return self._exec_plan + + def set_compiled(self, compiled: _CompiledRuntime) -> None: + self._compiled = compiled + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + residual: torch.Tensor | None, + aux_hidden_states: list | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + return self._compiled.forward( + positions, hidden_states, ctx, out_cache_loc, residual, aux_hidden_states + ) + + +class CompiledMoEDecoderLayer(CompiledDecoderLayer): + + @property + def is_moe_layer(self) -> bool: + return True diff --git a/python/tokenspeed/runtime/models/base/execution.py b/python/tokenspeed/runtime/models/base/execution.py new file mode 100644 index 0000000..33984c9 --- /dev/null +++ b/python/tokenspeed/runtime/models/base/execution.py @@ -0,0 +1,182 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +import torch +from torch import nn + +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.execution.context import ForwardContext + +if TYPE_CHECKING: + from tokenspeed.runtime.models.base.comm_ops import CommOp + +from tokenspeed.runtime.models.base.module_spec import ModuleKind, ModuleSpec +from tokenspeed.runtime.models.base.placement import Placement + + +@dataclass(frozen=True, slots=True) +class ExecutionNode: + module: nn.Module + spec: ModuleSpec + name: str | None = None + + +@dataclass(frozen=True, slots=True) +class ExecutionState: + hidden_states: torch.Tensor + residual: torch.Tensor | None + ctx: ForwardContext + out_cache_loc: torch.Tensor + + +StepRunner = Callable[[ExecutionState, torch.Tensor], ExecutionState] + + +@dataclass +class ExecutionStep: + runner: StepRunner + module: nn.Module | None = None + pre_comms: list[CommOp] = field(default_factory=list) + post_comms: list[CommOp] = field(default_factory=list) + spec: ModuleSpec = field(default_factory=ModuleSpec) + kind: ModuleKind = ModuleKind.GENERIC + captures_aux: bool = False + skip_on_idle: bool = False + name: str | None = None + + +class CompiledDecoderLayer(nn.Module): + def __init__( + self, + steps: list[ExecutionStep], + final_placement: Placement | None, + mapping: Mapping, + ) -> None: + from tokenspeed.runtime.models.base.comm_ops import ( + AllGatherOp, + ReduceScatterOp, + ResidualAllGatherOp, + ResidualSliceOp, + ) + + super().__init__() + self.final_placement = final_placement + self.mapping = mapping + + self.steps = steps + self.comm_modules = nn.ModuleList() + has_rsag_comms = False + for step in steps: + for comm in step.pre_comms: + self.comm_modules.append(comm) + if isinstance( + comm, + ( + AllGatherOp, + ReduceScatterOp, + ResidualAllGatherOp, + ResidualSliceOp, + ), + ): + has_rsag_comms = True + for comm in step.post_comms: + self.comm_modules.append(comm) + if isinstance( + comm, + ( + AllGatherOp, + ReduceScatterOp, + ResidualAllGatherOp, + ResidualSliceOp, + ), + ): + has_rsag_comms = True + self.has_rsag_comms = has_rsag_comms + + def can_fuse_embed_reduce(self, num_tokens: int) -> bool: + from tokenspeed.runtime.models.base.comm_ops import FusedReduceNormOp + + if not self.steps: + return False + first_module = self.steps[0].module + if isinstance(first_module, FusedReduceNormOp): + return first_module._should_fuse(num_tokens) + return False + + def _num_global_tokens(self, ctx: ForwardContext) -> int: + if ctx.global_num_tokens is not None: + return sum(ctx.global_num_tokens) + return ctx.input_num_tokens + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + residual: torch.Tensor | None, + aux_hidden_states: list | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + num_global_tokens = self._num_global_tokens(ctx) + is_idle = ctx.forward_mode.is_idle() if ctx.forward_mode else False + + if num_global_tokens == 0: + return hidden_states, residual + if hidden_states.shape[0] == 0 and not self.has_rsag_comms: + return hidden_states, residual + + state = ExecutionState(hidden_states, residual, ctx, out_cache_loc) + + for step in self.steps: + if is_idle and step.skip_on_idle: + continue + + for comm in step.pre_comms: + hidden_states, residual = comm( + state.hidden_states, state.residual, state.ctx + ) + state = ExecutionState( + hidden_states, residual, state.ctx, state.out_cache_loc + ) + + state = step.runner(state, positions) + + if ( + step.captures_aux + and aux_hidden_states is not None + and state.residual is not None + ): + aux_hidden_states.append(state.residual.clone()) + + for comm in step.post_comms: + hidden_states, residual = comm( + state.hidden_states, state.residual, state.ctx + ) + state = ExecutionState( + hidden_states, residual, state.ctx, state.out_cache_loc + ) + + return state.hidden_states, state.residual diff --git a/python/tokenspeed/runtime/models/base/module_spec.py b/python/tokenspeed/runtime/models/base/module_spec.py new file mode 100644 index 0000000..54aa3ac --- /dev/null +++ b/python/tokenspeed/runtime/models/base/module_spec.py @@ -0,0 +1,93 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum, auto + +from tokenspeed.runtime.models.base.placement import Placement + + +class ModuleKind(Enum): + NORM = auto() + ATTENTION = auto() + DENSE_MLP = auto() + MOE = auto() + GENERIC = auto() + + +class CallConvention(Enum): + NORM_WITH_OPTIONAL_RESIDUAL = auto() + ATTENTION = auto() + MOE = auto() + HIDDEN_STATES_ONLY = auto() + + +class FusionCapability(Enum): + NONE = auto() + REDUCE_NORM = auto() + + +@dataclass(frozen=True, slots=True) +class ModuleSpec: + input_placement: Placement | None = None + output_placement: Placement | None = None + kind: ModuleKind = ModuleKind.GENERIC + call: CallConvention = CallConvention.HIDDEN_STATES_ONLY + fusion: FusionCapability = FusionCapability.NONE + captures_aux: bool = False + skip_on_idle: bool = False + + @classmethod + def from_kind( + cls, + *, + input_placement: Placement | None = None, + output_placement: Placement | None = None, + kind: ModuleKind = ModuleKind.GENERIC, + supports_fused_reduce_norm: bool = False, + captures_aux: bool = False, + skip_on_idle: bool = False, + ) -> ModuleSpec: + if kind == ModuleKind.NORM: + call = CallConvention.NORM_WITH_OPTIONAL_RESIDUAL + elif kind == ModuleKind.ATTENTION: + call = CallConvention.ATTENTION + elif kind == ModuleKind.MOE: + call = CallConvention.MOE + else: + call = CallConvention.HIDDEN_STATES_ONLY + + fusion = ( + FusionCapability.REDUCE_NORM + if supports_fused_reduce_norm + else FusionCapability.NONE + ) + + return cls( + input_placement=input_placement, + output_placement=output_placement, + kind=kind, + call=call, + fusion=fusion, + captures_aux=captures_aux, + skip_on_idle=skip_on_idle, + ) diff --git a/python/tokenspeed/runtime/models/base/placement.py b/python/tokenspeed/runtime/models/base/placement.py new file mode 100644 index 0000000..4a849ad --- /dev/null +++ b/python/tokenspeed/runtime/models/base/placement.py @@ -0,0 +1,113 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""DTensor-inspired Placement type system for describing tensor distribution. + +Instead of using torch.DTensor directly (designed for training, high overhead), +we use lightweight annotations + compile-time static analysis to describe how +tensors are distributed across parallel groups. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum, auto + +from tokenspeed.runtime.distributed.mapping import Mapping + + +class PlacementType(Enum): + """Describes the distribution state of a tensor on a parallel dimension.""" + + REPLICATE = auto() # Every rank holds a full copy + SHARD = auto() # Scattered along the token dimension across ranks + PARTIAL = auto() # Each rank holds a partial sum; needs reduce to be complete + + +class ParallelGroup(Enum): + """The parallel group a communication belongs to.""" + + ATTN_TP = auto() + DENSE_TP = auto() + MOE_TP_EP = auto() + + +@dataclass(frozen=True, slots=True) +class Placement: + """Describes the distribution state of a tensor.""" + + type: PlacementType + group: ParallelGroup + + def __repr__(self) -> str: + return f"Placement({self.type.name}, {self.group.name})" + + @classmethod + def replicate(cls, group: ParallelGroup) -> Placement: + return cls(PlacementType.REPLICATE, group) + + @classmethod + def shard(cls, group: ParallelGroup) -> Placement: + return cls(PlacementType.SHARD, group) + + @classmethod + def partial(cls, group: ParallelGroup) -> Placement: + return cls(PlacementType.PARTIAL, group) + + +Replicate = Placement.replicate +Shard = Placement.shard +Partial = Placement.partial + + +def group_tp_size(mapping: Mapping, group: ParallelGroup) -> int: + if group == ParallelGroup.ATTN_TP: + return mapping.attn.tp_size + elif group == ParallelGroup.DENSE_TP: + return mapping.dense.tp_size + elif group == ParallelGroup.MOE_TP_EP: + return mapping.moe.tp_ep_size + else: + raise ValueError(f"Unknown group: {group}") + + +def group_has_parallel(mapping: Mapping, group: ParallelGroup) -> bool: + if group == ParallelGroup.ATTN_TP: + return mapping.has_attn_tp + elif group == ParallelGroup.DENSE_TP: + return mapping.dense.has_tp + elif group == ParallelGroup.MOE_TP_EP: + return mapping.moe.has_tp_ep + else: + raise ValueError(f"Unknown group: {group}") + + +def use_all_reduce( + mapping: Mapping, src_group: ParallelGroup, dst_group: ParallelGroup +) -> bool: + return group_tp_size(mapping, src_group) == group_tp_size(mapping, dst_group) + + +def can_fuse_reduce_norm( + mapping: Mapping, + src_group: ParallelGroup, + dst_group: ParallelGroup, +) -> bool: + return use_all_reduce(mapping, src_group, dst_group) and mapping.has_attn_tp diff --git a/python/tokenspeed/runtime/models/base/transformer_model.py b/python/tokenspeed/runtime/models/base/transformer_model.py new file mode 100644 index 0000000..c822c58 --- /dev/null +++ b/python/tokenspeed/runtime/models/base/transformer_model.py @@ -0,0 +1,242 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Base transformer model: embed -> layers -> norm.""" + +from __future__ import annotations + +import torch +from torch import nn +from transformers import PretrainedConfig + +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.layers.layernorm import RMSNorm +from tokenspeed.runtime.layers.quantization import QuantizationConfig +from tokenspeed.runtime.layers.vocab_parallel_embedding import VocabParallelEmbedding +from tokenspeed.runtime.models.base.comm_ops import FinalNormOp +from tokenspeed.runtime.models.base.compiler import ( + compile_decoder_layer, + find_first_compute_input_group, +) +from tokenspeed.runtime.models.base.decoder_layer import ( + BaseDecoderLayer, + CompiledDecoderLayer, +) +from tokenspeed.runtime.models.base.placement import ParallelGroup, PlacementType +from tokenspeed.runtime.moe.distribution_recorder import ( + get_global_expert_distribution_recorder, +) +from tokenspeed.runtime.utils import add_prefix, make_layers + + +class BaseTransformerModel(nn.Module): + + layer_cls: type[BaseDecoderLayer] = BaseDecoderLayer + + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + + super().__init__() + + self.config = config + self.quant_config = quant_config + self.mapping = mapping + self.padding_idx: int | None = getattr(config, "pad_token_id", None) + self.vocab_size: int = config.vocab_size + + self.embed_tokens = self.resolve_embed(config, prefix) + self.layers = self.resolve_layers(config, quant_config, prefix) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.layers_to_capture: list[int] = [] + + self._compile_decoder_stack() + + # Build the final norm op that handles cross-layer communication + # after the last decoder layer (fused allreduce + norm, or separate + # norm + all-gather for RSAG mode). + self._final_norm_op = self._build_final_norm_op() + + def _compile_decoder_stack(self) -> None: + """Compile only ``CompiledDecoderLayer`` instances.""" + prev_output_group = None + for idx, layer in enumerate(self.layers): + if not isinstance(layer, CompiledDecoderLayer): + continue + next_layer_input_group = None + if idx + 1 < len(self.layers): + next_layer = self.layers[idx + 1] + if isinstance(next_layer, CompiledDecoderLayer): + next_exec_plan = next_layer.resolve_exec_plan() + next_layer_input_group = find_first_compute_input_group( + next_exec_plan + ) + compiled = compile_decoder_layer( + layer=layer, + exec_plan=layer.resolve_exec_plan(), + mapping=self.mapping, + prev_layer_output_group=prev_output_group, + next_layer_input_group=next_layer_input_group, + ) + layer.set_compiled(compiled) + if compiled.final_placement is not None: + prev_output_group = compiled.final_placement.group + else: + prev_output_group = None + + def _build_final_norm_op(self) -> FinalNormOp: + """Create a FinalNormOp for the post-last-layer norm + comm.""" + last_layer = self.layers[-1] if len(self.layers) > 0 else None + + use_ar = True + group_type = ParallelGroup.ATTN_TP + if isinstance(last_layer, CompiledDecoderLayer): + compiled = getattr(last_layer, "_compiled", None) + if compiled is not None and compiled.final_placement is not None: + use_ar = compiled.final_placement.type != PlacementType.SHARD + group_type = compiled.final_placement.group + + return FinalNormOp( + mapping=self.mapping, + group_type=group_type, + norm_module=self.norm, + use_all_reduce_mode=use_ar, + lm_head_group_type=ParallelGroup.ATTN_TP, + ) + + def resolve_embed(self, config: PretrainedConfig, prefix: str) -> nn.Module: + return VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + prefix=add_prefix("embed_tokens", prefix), + ) + + def resolve_layers( + self, + config: PretrainedConfig, + quant_config: QuantizationConfig | None, + prefix: str, + ) -> nn.ModuleList: + + layer_cls = self.layer_cls + mapping = self.mapping + + return make_layers( + config.num_hidden_layers, + lambda idx, prefix: layer_cls( + config=config, + layer_id=idx, + mapping=mapping, + quant_config=quant_config, + prefix=prefix, + ), + prefix=add_prefix("layers", prefix), + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + input_embeds: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, list[torch.Tensor] | None]: + + hidden_states = input_embeds + residual = None + + if input_embeds is None: + # When TP > 1 and fused allreduce+norm is available, skip the + # NCCL allreduce in the embedding and let the first decoder layer + # fuse it with the input layernorm via the fused all-reduce kernel. + first_layer = self.layers[0] + if isinstance(first_layer, CompiledDecoderLayer): + first_compiled = first_layer._compiled + fuse_embed_reduce = first_compiled.can_fuse_embed_reduce( + input_ids.shape[0] + ) + elif isinstance(first_layer, BaseDecoderLayer): + fuse_embed_reduce = ( + self.mapping.attn.tp_size > 1 + and first_layer.comm_manager.should_fuse(input_ids.shape[0]) + ) + else: + fuse_embed_reduce = False + hidden_states = self.embed_tokens( + input_ids, reduce_results=not fuse_embed_reduce + ) + if fuse_embed_reduce: + residual = torch.zeros_like(hidden_states) + + aux_hidden_states: list[torch.Tensor] = [] + + for i, layer in enumerate(self.layers): + + with get_global_expert_distribution_recorder().with_current_layer(i): + + hidden_states, residual = layer( + positions, + hidden_states, + ctx, + out_cache_loc, + residual, + aux_hidden_states=( + aux_hidden_states if i in self.layers_to_capture else None + ), + ) + + if not ctx.forward_mode.is_idle(): + + if residual is None: + raise RuntimeError("residual is required for non-idle forward mode.") + + if isinstance(layer, BaseDecoderLayer): + hidden_states, final_residual = layer.comm_manager.final_norm( + hidden_states, residual, ctx, self.norm + ) + else: + hidden_states, final_residual = self._final_norm_op( + hidden_states, residual, ctx + ) + + # An id == num_layers (capture index num_layers + 1) selects the + # final norm's output residual as an aux state, matching how each + # layer type captures in-loop: BaseDecoderLayer gathers across + # attn-TP, CompiledDecoderLayer appends raw. + if ( + aux_hidden_states is not None + and final_residual is not None + and len(self.layers) + 1 in self.layers_to_capture + ): + if hasattr(layer, "comm_manager"): + final_residual = layer.comm_manager.gather_residual( + final_residual, ctx + ) + aux_hidden_states.append(final_residual.clone()) + + return hidden_states, aux_hidden_states or None diff --git a/python/tokenspeed/runtime/models/deepseek_nextn.py b/python/tokenspeed/runtime/models/deepseek_nextn.py new file mode 100755 index 0000000..b9fa7f3 --- /dev/null +++ b/python/tokenspeed/runtime/models/deepseek_nextn.py @@ -0,0 +1,488 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Inference-only DeepSeek NextN Speculative Decoding.""" + +from __future__ import annotations + +import logging +from collections.abc import Iterable + +import torch +from torch import nn +from transformers import PretrainedConfig + +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.execution.context import ( + ForwardContext, + report_collective_sizing, +) +from tokenspeed.runtime.layers.layernorm import RMSNorm +from tokenspeed.runtime.layers.linear import ReplicatedLinear +from tokenspeed.runtime.layers.logits_processor import LogitsMetadata, LogitsProcessor +from tokenspeed.runtime.layers.moe import ( + ExpertCheckpointSchema, + build_moe_checkpoint_loader, +) +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.layers.quantization.utils import block_dequant +from tokenspeed.runtime.layers.utils import ( + CP_METADATA, + ENABLE_CP, + cp_all_gather_rerange_output, + cp_split_and_rebuild_data, +) +from tokenspeed.runtime.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader +from tokenspeed.runtime.models.deepseek_v3 import ( + DeepseekV3DecoderLayer, + DeepseekV3DraftAttentionMLA, + DeepseekV3ForCausalLM, +) + +logger = logging.getLogger(__name__) + + +class DeepseekV3DraftDecoderLayer(DeepseekV3DecoderLayer): + """Decoder layer that injects the draft attention and narrows residuals. + + Restricted to single-layer drafts: ``_apply_correction`` mutates + ``ctx.draft_seq_lens_buf`` in place and is not idempotent across layers. + """ + + @property + def attention_cls(self) -> type[nn.Module]: + return DeepseekV3DraftAttentionMLA + + def _maybe_narrow_residual( + self, + residual: torch.Tensor, + ctx: ForwardContext, + ) -> torch.Tensor: + """Narrow residual to the draft attention's [bs, H] live rows.""" + if ctx.accept_lengths is None or ctx.forward_mode.is_idle(): + return residual + return residual.index_select(0, ctx.gather_ids) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + residual: torch.Tensor | None, + ) -> torch.Tensor: + num_global_tokens, max_num_tokens_per_gpu = self.comm_manager.get_num_tokens( + ctx + ) + + if not ctx.forward_mode.is_idle(): + hidden_states, residual = self.comm_manager.input_reduce_norm( + hidden_states, residual + ) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ctx=ctx, + out_cache_loc=out_cache_loc, + comm_manager=self.comm_manager, + ) + residual = self._maybe_narrow_residual(residual, ctx) + hidden_states, residual = self.comm_manager.post_attn_reduce_norm( + hidden_states, residual, ctx + ) + hidden_states = self.forward_mlp( + hidden_states, + residual, + ctx, + num_global_tokens, + max_num_tokens_per_gpu, + ) + else: + hidden_states = self.forward_mlp( + hidden_states, + residual, + ctx, + num_global_tokens, + max_num_tokens_per_gpu, + ) + return hidden_states, residual + + +class DeepseekModelNextN(nn.Module): + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + ) -> None: + super().__init__() + self.mapping = mapping + self.vocab_size = config.vocab_size + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.eh_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False) + + self.alt_stream = torch.cuda.Stream() + self.decoder = DeepseekV3DraftDecoderLayer( + config, + 0, + mapping=self.mapping, + quant_config=quant_config, + is_nextn=True, + alt_stream=self.alt_stream, + ) + + self.shared_head = nn.Module() + self.shared_head.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + input_embeds: torch.Tensor | None = None, + captured_hidden_states: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + if captured_hidden_states is None: + raise ValueError("DeepSeek NextN requires captured_hidden_states.") + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + + hidden_states = self.eh_proj( + torch.cat( + ( + self.enorm(hidden_states), + self.hnorm(captured_hidden_states), + ), + dim=-1, + ) + ) + + residual = None + if CP_METADATA: + hidden_states = cp_split_and_rebuild_data( + hidden_states, + CP_METADATA.value.split_list, + CP_METADATA.value.zigzag_index, + ) + positions = cp_split_and_rebuild_data( + positions, CP_METADATA.value.split_list, CP_METADATA.value.zigzag_index + ) + hidden_states, residual = self.decoder( + positions, + hidden_states, + ctx, + out_cache_loc, + residual, + ) + + if not ctx.forward_mode.is_idle(): + if not ENABLE_CP: + hidden_states, _ = self.decoder.comm_manager.final_norm( + hidden_states, residual, ctx, self.shared_head.norm + ) + else: + hidden_states, _ = self.shared_head.norm(hidden_states, residual) + if CP_METADATA: + hidden_states = cp_all_gather_rerange_output( + hidden_states, + CP_METADATA.value, + self.mapping.attn.tp_rank, + self.mapping.attn.tp_group, + ) + return hidden_states, None + + +class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM): + + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + ) -> None: + nn.Module.__init__(self) + self.config = config + self.mapping = mapping + + # FP4 quantization is not used for the NextN draft model. + # The NVIDIA FP4 checkpoint stores NextN MoE weights in BF16, + # so the draft model runs entirely in BF16. + if quant_config is not None and quant_config.get_name() == "nvfp4": + logger.warning( + "Overriding DeepseekV3ForCausalLMNextN quant config: " + "FP4 quantization not used for NextN draft model." + ) + quant_config = None + + self.quant_config = quant_config + + self.model = DeepseekModelNextN( + config, mapping=self.mapping, quant_config=quant_config + ) + + if self.mapping.attn.has_dp: + self.lm_head = ReplicatedLinear( + config.hidden_size, + config.vocab_size, + bias=False, + ) + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + self.logits_processor = LogitsProcessor( + config, + skip_all_gather=self.mapping.attn.has_dp, + do_argmax=True, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + + @torch.no_grad() + def forward( + self, + ctx: ForwardContext, + input_ids: torch.Tensor, + positions: torch.Tensor, + out_cache_loc: torch.Tensor, + captured_hidden_states: torch.Tensor | None = None, + ) -> torch.Tensor: + with report_collective_sizing(ctx, ctx.bs, ctx.global_bs): + hidden_states, _ = self.model( + input_ids, + positions, + ctx, + out_cache_loc, + captured_hidden_states=captured_hidden_states, + ) + logits_metadata = LogitsMetadata.from_forward_context(ctx) + return self.logits_processor( + input_ids, hidden_states, self.lm_head, logits_metadata + ) + + def get_hot_token_id(self): + # MTP drafts every vocab token; the hot-token-id mechanism is an + # EAGLE3-only optimization (see deepseek_v3.py:2063, llama_eagle3.py). + return None + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + # Fuse q_a_proj and kv_a_proj_with_mqa along output dimension when q_lora_rank is not None + fuse_qkv_a_proj = hasattr(self.config, "q_lora_rank") and ( + self.config.q_lora_rank is not None + ) + cached_a_proj = {} if fuse_qkv_a_proj else None + + nextn_spec_weight_names = [ + "shared_head.norm", + "eh_proj", + "enorm", + "hnorm", + ] + + params_dict = dict(self.named_parameters()) + # MoE expert weights, scales, and activation scales are handled + # by the checkpoint loader. + moe_loader = build_moe_checkpoint_loader( + params_dict=params_dict, + expert_schema=ExpertCheckpointSchema( + gate_proj_name="gate_proj", + down_proj_name="down_proj", + up_proj_name="up_proj", + ), + num_experts=self.config.n_routed_experts, + ep_rank=self.mapping.moe.ep_rank, + ep_size=self.mapping.moe.ep_size, + ) + for name, loaded_weight in weights: + if hasattr(self.config, "num_nextn_predict_layers"): + num_nextn_layers = self.config.num_nextn_predict_layers + if num_nextn_layers != 1: + raise ValueError("Only 1 nextn layer is supported") + nextn_layer_prefix = "model.layers.0" + if num_nextn_layers != self.config.num_hidden_layers: + if name.startswith("model.layers"): + name_list = name.split(".") + if ( + len(name_list) >= 3 + and int(name_list[2]) >= self.config.num_hidden_layers + ): + nextn_layer_prefix = "model.layers." + str(name_list[2]) + else: + continue + if not name.startswith(nextn_layer_prefix): + continue + else: + raise ValueError("num_nextn_predict_layers is not in the config") + + # Use shared head and embed weights from target model + if "shared_head.head" in name or "embed_tokens" in name: + continue + + is_decoder = True + # For nextn specific weights + for weight_name in nextn_spec_weight_names: + if weight_name in name: + name = name.replace(nextn_layer_prefix, "model") + is_decoder = False + break + # For decoder layer weights + if is_decoder: + name = name.replace(nextn_layer_prefix, "model.decoder") + + if "rotary_emb.inv_freq" in name: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + # Skip non-stacked layers and experts (experts handled below). + if weight_name not in name: + continue + # We have mlp.experts[0].gate_proj in the checkpoint. + # Since moe_loader handles the experts below, + # we need to skip here BEFORE we update the name, otherwise + # name will be updated to mlp.experts[0].gate_up_proj, which + # will then be updated below by moe_loader + # for mlp.experts[0].gate_gate_up_proj, which breaks load. + if ("mlp.experts." in name) and name not in params_dict: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if moe_loader.matches(name): + moe_loader.load(name, loaded_weight) + continue + + if fuse_qkv_a_proj and ( + "q_a_proj" in name or "kv_a_proj_with_mqa" in name + ): + cached_a_proj[name] = loaded_weight + q_a_proj_name = ( + name + if "q_a_proj" in name + else name.replace("kv_a_proj_with_mqa", "q_a_proj") + ) + kv_a_proj_name = ( + name + if "kv_a_proj_with_mqa" in name + else name.replace("q_a_proj", "kv_a_proj_with_mqa") + ) + + # When both q_a_proj and kv_a_proj_with_mqa has been cached, load the fused weight to parameter + if ( + q_a_proj_name in cached_a_proj + and kv_a_proj_name in cached_a_proj + ): + + q_a_proj_weight = cached_a_proj[q_a_proj_name] + kv_a_proj_weight = cached_a_proj[kv_a_proj_name] + fused_weight = torch.cat( + [q_a_proj_weight, kv_a_proj_weight], dim=0 + ) + + if "q_a_proj" in name: + param_name = name.replace( + "q_a_proj", "fused_qkv_a_proj_with_mqa" + ) + else: + param_name = name.replace( + "kv_a_proj_with_mqa", "fused_qkv_a_proj_with_mqa" + ) + param = params_dict[param_name] + + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, fused_weight) + cached_a_proj.pop(q_a_proj_name) + cached_a_proj.pop(kv_a_proj_name) + else: + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + self.post_load_weights() + + def post_load_weights(self): + self_attn = self.model.decoder.self_attn + if ( + hasattr(self.quant_config, "weight_block_size") + and (self.quant_config.weight_block_size is not None) + and self_attn.kv_b_proj.weight.dtype + in ( + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + ) + ): + weight_block_size = self.quant_config.weight_block_size + dtype = torch.get_default_dtype() + w = block_dequant( + self_attn.kv_b_proj.weight, + self_attn.kv_b_proj.weight_scale_inv, + weight_block_size, + ).to(dtype) + else: + w = self_attn.kv_b_proj.weight + + 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) + self_attn.w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2) + self_attn.w_vc = w_vc.contiguous().transpose(1, 2) + + +EntryClass = [DeepseekV3ForCausalLMNextN] diff --git a/python/tokenspeed/runtime/models/deepseek_v3.py b/python/tokenspeed/runtime/models/deepseek_v3.py new file mode 100644 index 0000000..a0d87a1 --- /dev/null +++ b/python/tokenspeed/runtime/models/deepseek_v3.py @@ -0,0 +1,2198 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Inference-only DeepseekV3 model.""" + +# ruff: noqa: E402 + +from __future__ import annotations + +import re +from collections.abc import Iterable +from dataclasses import replace +from typing import Any + +import torch +import torch.nn.functional as F +from tokenspeed_kernel.ops.attention import attn_merge_state +from tokenspeed_kernel.ops.attention.tokenspeed_mla import mla_kv_pack_quantize_fp8 +from tokenspeed_kernel.ops.embedding import apply_rope_mla +from tokenspeed_kernel.ops.gemm.cute_dsl import ( + nvfp4_gemm_swiglu_nvfp4_quant, +) +from tokenspeed_kernel.ops.gemm.trtllm import dsv3_fused_a_gemm +from tokenspeed_kernel.ops.quantization.flashinfer import fp4_quantize +from tokenspeed_kernel.ops.quantization.triton import fp8_quantize +from tokenspeed_kernel.platform import current_platform +from tokenspeed_kernel.thirdparty.cuda import dsv3_router_gemm, moe_finalize_fuse_shared +from torch import nn +from transformers import PretrainedConfig + +from tokenspeed.runtime.configs.utils import get_rope_theta +from tokenspeed.runtime.layers.moe import ( + ExpertCheckpointSchema, + build_moe_checkpoint_loader, +) +from tokenspeed.runtime.layers.utils import ( + CP_METADATA, + ENABLE_CP, + cp_all_gather_rerange_output, + cp_split_and_rebuild_data, + get_layer_id, +) + +_platform = current_platform() +_is_amd = _platform.is_amd +_is_blackwell = _platform.is_blackwell +_is_hopper_plus = _platform.is_hopper_plus +_device_sm = _platform.arch_version.major * 10 + _platform.arch_version.minor + +from tokenspeed.runtime.distributed import Mapping +from tokenspeed.runtime.distributed.comm_manager import CommManager +from tokenspeed.runtime.execution.breakable_cuda_graph import ( + break_point, + scrub_padding_tail, +) +from tokenspeed.runtime.execution.context import ( + ForwardContext, + report_collective_sizing, +) +from tokenspeed.runtime.execution.cuda_graph_wrapper import get_is_capture_mode +from tokenspeed.runtime.execution.forward_batch_info import ForwardMode +from tokenspeed.runtime.layers.activation import SiluAndMul +from tokenspeed.runtime.layers.dense.nvfp4 import Nvfp4LinearMethod +from tokenspeed.runtime.layers.layernorm import FusedRMSNorm, RMSNorm +from tokenspeed.runtime.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from tokenspeed.runtime.layers.logits_processor import LogitsProcessor +from tokenspeed.runtime.layers.moe.expert import MoELayer +from tokenspeed.runtime.layers.moe.topk import TopK +from tokenspeed.runtime.layers.moe.utils import RoutingMethodType +from tokenspeed.runtime.layers.paged_attention import PagedAttention +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.layers.quantization.nvfp4 import Nvfp4Config +from tokenspeed.runtime.layers.quantization.utils import ( + block_dequant, + should_exclude_quant_module, +) +from tokenspeed.runtime.layers.rotary_embedding import get_rope +from tokenspeed.runtime.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from tokenspeed.runtime.model_loader.weight_utils import ( + default_weight_loader, + kv_cache_scales_loader, +) +from tokenspeed.runtime.models.base import BaseCausalLM +from tokenspeed.runtime.models.utils import ( + create_fused_set_kv_buffer_arg, +) +from tokenspeed.runtime.moe.distribution_recorder import ( + get_global_expert_distribution_recorder, +) +from tokenspeed.runtime.moe.expert_location import ModelConfigForExpertLocation +from tokenspeed.runtime.utils import LazyValue, add_prefix, get_colorful_logger +from tokenspeed.runtime.utils.cuda_stream import StreamFork +from tokenspeed.runtime.utils.env import envs, global_server_args_dict +from tokenspeed.runtime.utils.pdl import pdl_enabled + +logger = get_colorful_logger(__name__) + +_OPTIONAL_MISSING_WEIGHT_SUFFIXES = ( + ".k_scale", + ".v_scale", +) + + +def _prepare_mla_kv_b_proj_weights( + w: torch.Tensor, self_attn +) -> tuple[torch.Tensor, torch.Tensor]: + 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) + if _is_amd: + return w_kc.contiguous(), w_vc.transpose(1, 2).contiguous() + return ( + w_kc.transpose(1, 2).contiguous().transpose(1, 2), + w_vc.contiguous().transpose(1, 2), + ) + + +class DeepseekV3MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + is_shared_expert: bool = False, + ) -> None: + super().__init__() + self.mapping = mapping + if is_shared_expert: + tp_rank = self.mapping.moe.tp_ep_rank + tp_size = self.mapping.moe.tp_ep_size + tp_group = self.mapping.moe.tp_ep_group + else: + tp_rank = self.mapping.dense.tp_rank + tp_size = self.mapping.dense.tp_size + tp_group = self.mapping.dense.tp_group + + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + tp_size=tp_size, + tp_rank=tp_rank, + tp_group=tp_group, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + reduce_results=False, # Communication is handled externally and manually controlled + tp_size=tp_size, + tp_rank=tp_rank, + tp_group=tp_group, + quant_config=quant_config, + prefix=add_prefix("down_proj", prefix), + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. Only silu is supported for now." + ) + self.act_fn = SiluAndMul() + self._use_nvfp4_gemm_swiglu_nvfp4_quant = ( + envs.TOKENSPEED_NVFP4_GEMM_SWIGLU_NVFP4_QUANT.get() + and _is_blackwell + and isinstance(self.gate_up_proj.quant_method, Nvfp4LinearMethod) + and isinstance(self.down_proj.quant_method, Nvfp4LinearMethod) + ) + self.gate_up_proj.interleave_linear_and_gate = ( + self._use_nvfp4_gemm_swiglu_nvfp4_quant + ) + + def forward(self, x): + if x.size(0) == 0: + return x + + if self._use_nvfp4_gemm_swiglu_nvfp4_quant: + x_fc1_fp4, x_fc1_scale = fp4_quantize( + x, self.gate_up_proj.input_scale_inv, enable_pdl=pdl_enabled() + ) + x_fp4, x_scale = nvfp4_gemm_swiglu_nvfp4_quant( + x_fc1_fp4, + x_fc1_scale, + self.gate_up_proj.weight_swiglu_interleaved, + self.gate_up_proj.weight_scale_swiglu_interleaved, + self.gate_up_proj.alpha, + self.down_proj.input_scale_inv, + enable_pdl=pdl_enabled(), + ) + x, _ = self.down_proj((x_fp4, x_scale)) + return x + + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class MoEGate(nn.Module): + _DSV3_ROUTER_GEMM_EXPERTS = (256, 384, 768) + _DSV3_ROUTER_GEMM_HIDDEN = (3072, 6144, 7168) + + def __init__(self, config, prefix: str = ""): + super().__init__() + self.weight = nn.Parameter( + torch.empty((config.n_routed_experts, config.hidden_size)) + ) + if config.topk_method == "noaux_tc": + self.e_score_correction_bias = nn.Parameter( + torch.empty((config.n_routed_experts), dtype=torch.float32) + ) + else: + self.e_score_correction_bias = None + + self.use_dsv3_router_gemm = ( + _is_hopper_plus + and self.weight.dtype in (torch.bfloat16, torch.float32) + and config.n_routed_experts in self._DSV3_ROUTER_GEMM_EXPERTS + and config.hidden_size in self._DSV3_ROUTER_GEMM_HIDDEN + ) + + def forward(self, hidden_states, comm_manager=None): + if self.use_dsv3_router_gemm and hidden_states.size(0) > 0: + logits = dsv3_router_gemm( + hidden_states, + self.weight, + out_dtype=torch.float32, + enable_pdl=pdl_enabled(), + ) + else: + logits = F.linear(hidden_states, self.weight, None) + return logits + + +class DeepseekV3MoE(nn.Module): + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + layer_index: int = -1, + prefix: str = "", + alt_stream: torch.cuda.Stream | None = None, + ): + super().__init__() + self.mapping = mapping + self.layer_index = layer_index + self.n_shared_experts = config.n_shared_experts + self.routed_scaling_factor = config.routed_scaling_factor + self.stream_fork = StreamFork(alt_stream) + + if self.mapping.moe.ep_size > config.n_routed_experts: + raise ValueError( + f"EP size {self.mapping.moe.ep_size} is greater than the number of experts {config.n_routed_experts}." + ) + if config.hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {config.hidden_act}. Only silu is supported for now." + ) + + self.gate = MoEGate(config=config, prefix=add_prefix("gate", prefix)) + + if config.n_shared_experts is not None: + intermediate_size = config.moe_intermediate_size * config.n_shared_experts + self.shared_experts = DeepseekV3MLP( + hidden_size=config.hidden_size, + intermediate_size=intermediate_size, + hidden_act=config.hidden_act, + mapping=self.mapping, + quant_config=quant_config, + prefix=add_prefix("shared_experts", prefix), + is_shared_expert=True, + ) + + self.experts = MoELayer( + top_k=config.num_experts_per_tok, + num_experts=config.n_routed_experts + + global_server_args_dict["ep_num_redundant_experts"], + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + quant_config=quant_config, + layer_index=layer_index, + prefix=prefix, + tp_rank=self.mapping.moe.tp_rank, + tp_size=self.mapping.moe.tp_size, + ep_rank=self.mapping.moe.ep_rank, + ep_size=self.mapping.moe.ep_size, + routing_config={ + "n_group": getattr(config, "n_group", 0), + "topk_group": getattr(config, "topk_group", 0), + "routed_scaling_factor": getattr(config, "routed_scaling_factor", 1.0), + "normalize_topk_weights": config.norm_topk_prob, + "correction_bias": self.gate.e_score_correction_bias, + "routing_method_type": RoutingMethodType.DeepSeekV3, + }, + ) + + self.topk = TopK( + top_k=config.num_experts_per_tok, + renormalize=config.norm_topk_prob, + use_grouped_topk=True, + num_expert_group=config.n_group, + num_fused_shared_experts=0, + topk_group=config.topk_group, + correction_bias=self.gate.e_score_correction_bias, + routed_scaling_factor=self.routed_scaling_factor, + output_format=self.experts.topk_output_format, + ) + + def get_moe_routed_weights(self): + return [ + x.data + for name, x in self.experts.named_parameters() + if name not in ["correction_bias"] and "shared_experts" not in name + ] + + def forward( + self, + hidden_states: torch.Tensor, + num_global_tokens: int, + max_num_tokens_per_gpu: int, + ) -> torch.Tensor: + num_tokens = hidden_states.size(0) + + with self.stream_fork.scope(enable=get_is_capture_mode()) as fork: + # router_logits: (num_tokens, n_experts) + router_logits = self.gate(hidden_states) + if num_tokens > 0: + topk_output = self.topk(hidden_states, router_logits) + else: + topk_output = self.topk.empty_topk_output( + hidden_states.device, + hidden_states=hidden_states, + router_logits=router_logits, + ) + + deferred_finalize = self.experts.supports_deferred_finalize + routed_expert_output = self.experts( + hidden_states=hidden_states, + topk_output=topk_output, + num_global_tokens=num_global_tokens, + max_num_tokens_per_gpu=max_num_tokens_per_gpu, + do_finalize=not deferred_finalize, + ) + + shared_output = None + with fork.branch(): + if self.n_shared_experts is not None and num_tokens > 0: + shared_output = self.shared_experts(hidden_states) + + if deferred_finalize: + gemm2_out, expert_weights, expanded_idx = routed_expert_output + final_hidden_states = moe_finalize_fuse_shared( + gemm2_out, + expanded_idx, + expert_weights, + shared_output, + top_k=self.topk.topk_config.top_k, + enable_pdl=pdl_enabled(), + ) + else: + final_hidden_states = ( + routed_expert_output + shared_output + if shared_output is not None + else routed_expert_output + ) + return final_hidden_states + + +def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float: + import math + + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + +class DeepseekV3FusedQkvAProjWithMqa(ReplicatedLinear): + def __init__( + self, + input_size: int, + output_size: int, + bias: bool = True, + skip_bias_add: bool = False, + params_dtype: torch.dtype | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + # ModelOpt NVFP4 checkpoints (e.g. DeepSeek-R1-0528-NVFP4-v2) keep the + # q_a_proj / kv_a_proj_with_mqa weights as bf16 via exclude_modules. + # exclude_modules matches by component name, not by the fused parent + # prefix, so the fused layer would otherwise allocate an NVFP4-packed + # buffer and crash when bf16 weights are copied in. + if isinstance(quant_config, Nvfp4Config) and prefix: + q_a_prefix = prefix.replace("fused_qkv_a_proj_with_mqa", "q_a_proj") + kv_a_prefix = prefix.replace( + "fused_qkv_a_proj_with_mqa", "kv_a_proj_with_mqa" + ) + if should_exclude_quant_module( + q_a_prefix, quant_config.exclude_modules + ) or should_exclude_quant_module(kv_a_prefix, quant_config.exclude_modules): + quant_config = None + super().__init__( + input_size, + output_size, + bias=bias, + skip_bias_add=skip_bias_add, + params_dtype=params_dtype, + quant_config=quant_config, + prefix=prefix, + ) + self.use_min_latency = ( + self.bias is None + and self.weight.dtype == torch.bfloat16 + and self.weight.size() == (2112, 7168) + and current_platform().is_nvidia + and _device_sm >= 90 + and _device_sm not in (120, 121) + ) + + def forward( + self, x: torch.Tensor, block_scale=None, output_dtype=None + ) -> torch.Tensor: + if ( + self.use_min_latency + and x.size(0) > 0 + and block_scale is None + and (output_dtype is None or output_dtype == torch.bfloat16) + ): + return dsv3_fused_a_gemm(x, self.weight.T) + + return super().forward(x, block_scale=block_scale, output_dtype=output_dtype)[0] + + +class DeepseekV3AttentionMLA(nn.Module): + # Backends that use non-absorbed MLA kernels (ragged prefill, paged KV decode). + _MLA_KERNEL_BACKENDS = ("mla", "trtllm_mla", "tokenspeed_mla") + # Backends that support chunked ragged prefill with prefix replay. + _RAGGED_PREFILL_BACKENDS = ("mla", "trtllm_mla", "tokenspeed_mla") + + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + hidden_size: int, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + q_lora_rank: int, + kv_lora_rank: int, + rope_theta: float = 10000, + rope_scaling: dict[str, Any] | None = None, + max_position_embeddings: int = 8192, + quant_config: QuantizationConfig | None = None, + layer_id=None, + prefix: str = "", + reduce_attn_results=True, + alt_stream: torch.cuda.Stream | None = None, + skip_rope: bool = False, + ) -> None: + super().__init__() + self.mapping = mapping + self.layer_id = layer_id + self.hidden_size = hidden_size + self.qk_nope_head_dim = qk_nope_head_dim + self.qk_rope_head_dim = qk_rope_head_dim + self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self.v_head_dim = v_head_dim + self.q_lora_rank = q_lora_rank + self.kv_lora_rank = kv_lora_rank + self.num_heads = num_heads + if num_heads % self.mapping.attn.tp_size != 0: + raise ValueError( + f"num_heads={num_heads} must be divisible by attn_tp_size={self.mapping.attn.tp_size}." + ) + self.num_local_heads = num_heads // self.mapping.attn.tp_size + self.scaling = self.qk_head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + self.config = config + self.alt_stream = alt_stream + self.attention_backend = global_server_args_dict["attention_backend"] + self.cli_factor = getattr(config, "cli_factor", 1) + self.prefix = prefix + + # modification to rope_scaling must be done early enough, b/c e.g. Indexer needs it + if rope_scaling: + rope_scaling["rope_type"] = "deepseek_yarn" + + if self.q_lora_rank is not None: + self.fused_qkv_a_proj_with_mqa = DeepseekV3FusedQkvAProjWithMqa( + self.hidden_size, + self.q_lora_rank + self.kv_lora_rank + self.qk_rope_head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("fused_qkv_a_proj_with_mqa", prefix), + ) + + self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) + self.q_b_proj = ColumnParallelLinear( + q_lora_rank, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("q_b_proj", prefix), + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + else: + self.q_proj = ColumnParallelLinear( + self.hidden_size, + self.num_heads * self.qk_head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("q_proj", prefix), + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + + self.kv_a_proj_with_mqa = ReplicatedLinear( + self.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("kv_a_proj_with_mqa", prefix), + ) + + self.kv_b_proj = ColumnParallelLinear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + quant_config=quant_config, + prefix=add_prefix("kv_b_proj", prefix), + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + # O projection. + self.o_proj = RowParallelLinear( + self.num_heads * self.v_head_dim, + self.hidden_size, + bias=False, + reduce_results=reduce_attn_results, + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps) + + # Fusion layer + if self.q_lora_rank is not None: + self.fused_qk_layernorm = FusedRMSNorm( + self.q_a_layernorm, + self.kv_a_layernorm, + ) + + if not skip_rope: + self.rotary_emb = get_rope( + qk_rope_head_dim, + rotary_dim=qk_rope_head_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + is_neox_style=False, + ) + + if rope_scaling: + mscale_all_dim = rope_scaling.get("mscale_all_dim", False) + scaling_factor = rope_scaling["factor"] + mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim)) + self.scaling = self.scaling * mscale * mscale + else: + self.rotary_emb = None + + # Fused RoPE+KV write kernel is incompatible with MLA: it assumes + # K and V have the same head_dim, but MLA's KV cache is a single + # [latent(512)|rope(64)] buffer where the two dimensions differ. + # Passing this to the kernel causes thread overflow and silent + # corruption of the latent cache. All DeepSeek V2/V3 models use + # MLA (kv_lora_rank > 0), so we unconditionally disable it here. + self.use_fused_set_kv_buffer = False + + self.attn_mqa = PagedAttention( + self.num_local_heads, + self.kv_lora_rank + self.qk_rope_head_dim, + self.scaling, + num_kv_heads=1, + layer_id=layer_id, + v_head_dim=self.kv_lora_rank, + ) + + self.attn_mha = PagedAttention( + self.num_local_heads, + self.qk_nope_head_dim + self.qk_rope_head_dim, + self.scaling, + num_kv_heads=self.num_local_heads, + layer_id=layer_id, + v_head_dim=self.v_head_dim, + ) + + self.w_kc = None + self.w_vc = None + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + comm_manager: CommManager, + block_scale: torch.Tensor | None = None, + ) -> torch.Tensor: + """MLA attention with a NARROW prefill-graph break. + + The token-shaped input/output projections (q/kv-down, layernorm, + q_b_proj, o_proj) stay in the captured prefill graph; only the + data-dependent attention -- KV write + varlen prefill / absorb decode + kernels + the live prefill/decode split -- runs as the eager break + (``_attn``). This keeps the big projection GEMMs graphed instead of + dispatch-bound eager, collapsing the inter-segment bubbles a coarse + whole-attention break leaves. Outside capture the ``@break_point`` is + a direct call, so the eager path is unchanged. + """ + if hidden_states.shape[0] == 0: + return hidden_states + q, latent_cache = self._project_q_latent( + hidden_states, ctx, comm_manager, block_scale + ) + attn_output = self._attn(positions, q, latent_cache, ctx, out_cache_loc) + output, _ = self.o_proj(attn_output) + return output + + def _project_q_latent( + self, + hidden_states: torch.Tensor, + ctx: ForwardContext, + comm_manager: CommManager, + block_scale: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """QKV projection producing absorbed ``q`` and raw ``latent_cache``.""" + if self.q_lora_rank is not None: + qkv = self.fused_qkv_a_proj_with_mqa( + hidden_states, block_scale, torch.bfloat16 + ) + qkv = comm_manager.pre_attn_comm(qkv, ctx) + q_a, latent_cache = qkv.split( + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], + dim=-1, + ) + kv_a = latent_cache[..., : self.kv_lora_rank] + q_norm = torch.empty_like(q_a) + if q_a.size(0) > 0: + self.fused_qk_layernorm( + input_q_a=q_a, input_kv_a=kv_a, output_q_a=q_norm + ) + q = self.q_b_proj(q_norm)[0] + else: + hidden_states = comm_manager.pre_attn_comm(hidden_states, ctx) + q = self.q_proj(hidden_states)[0] + latent_cache = self.kv_a_proj_with_mqa(hidden_states)[0] + kv_a = latent_cache[..., : self.kv_lora_rank] + self.kv_a_layernorm(kv_a, inplace=True) + return q, latent_cache + + @break_point + def _attn( + self, + positions: torch.Tensor, + q: torch.Tensor, + latent_cache: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + ) -> torch.Tensor: + """The eager break: KV write + varlen prefill / absorb decode attention. + + Prefill/decode dispatch over the full rows; subclasses override (the + draft variant narrows to live rows, see ``DeepseekV3DraftAttentionMLA``). + The split is recovered from LIVE state -- correct both in eager and + under a prefill-graph replay, where ``ctx`` is the live ambient context + but ``q`` is padded to the graph bucket (``q.size(0)`` is NOT the real + token count). The decode token count comes from the live ctx; the real + prefill token count from the live attention metadata (the same source + the padding scrub uses). Padded tail rows produce discarded garbage. + """ + spec = ctx.attn_backend.spec_num_tokens or 1 + num_decodes = max(ctx.bs - ctx.num_extends, 0) + num_decode_tokens = num_decodes * spec + if ctx.num_extends > 0: + cmeta = ctx.attn_backend.chunked_prefill_metadata + num_prefill_tokens = int(sum(cmeta.extend_seq_lens_cpu)) + else: + num_prefill_tokens = 0 + real_total = num_prefill_tokens + num_decode_tokens + attn_output = torch.empty( + q.size(0), + self.num_local_heads * self.v_head_dim, + dtype=q.dtype, + device=q.device, + ) + + if num_prefill_tokens > 0: + prefill_ctx = replace( + ctx, + bs=max(ctx.bs - num_decodes, 1), + num_extends=max(ctx.bs - num_decodes, 1), + input_num_tokens=num_prefill_tokens, + forward_mode=ForwardMode.EXTEND, + ) + self.forward_normal_chunked( + positions[:num_prefill_tokens], + q[:num_prefill_tokens], + latent_cache[:num_prefill_tokens], + prefill_ctx, + out_cache_loc[:num_prefill_tokens], + attn_output[:num_prefill_tokens], + ) + + if num_decode_tokens > 0: + decode_ctx = replace( + ctx, + bs=num_decodes, + num_extends=0, + input_num_tokens=num_decode_tokens, + forward_mode=ForwardMode.DECODE, + ) + self.forward_absorb( + positions[num_prefill_tokens:real_total], + q[num_prefill_tokens:real_total], + latent_cache[num_prefill_tokens:real_total], + decode_ctx, + out_cache_loc[num_prefill_tokens:real_total], + attn_output[num_prefill_tokens:real_total], + ) + + return attn_output + + def forward_absorb( + self, + positions: torch.Tensor, + q: torch.Tensor, + latent_cache: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + output: torch.Tensor, + ) -> torch.Tensor: + Q, K = self.forward_absorb_qkv_proj( + q, + latent_cache, + positions, + ctx, + out_cache_loc, + ) + return self.forward_absorb_attn_v_proj(Q, K, ctx, out_cache_loc, output) + + def forward_absorb_qkv_proj( + self, + q: torch.Tensor, + latent_cache: torch.Tensor, + positions, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + q = q.view(-1, self.num_local_heads, self.qk_head_dim) + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + + Q = torch.empty( + q_nope.size(0), + self.num_local_heads, + self.kv_lora_rank + self.qk_rope_head_dim, + dtype=q_nope.dtype, + device=q_nope.device, + ) + # latent_cache contains normalized kv_a and k_pe before rotate. + K = latent_cache.unsqueeze(1) + q_nope_out_view = Q[..., : self.kv_lora_rank] + if _is_amd: + q_nope_projected = torch.bmm( + q_nope.transpose(0, 1).contiguous(), + self.w_kc.contiguous(), + ) + q_nope_out_view.copy_(q_nope_projected.transpose(0, 1)) + else: + torch.bmm( + q_nope.transpose(0, 1), + self.w_kc, + out=q_nope_out_view.transpose(0, 1), + ) + # Model-owned fused FP8 decode: RoPE + quantize + KV cache write + # all done here, so backend only needs to do attention. + k_scale = getattr(self.attn_mqa, "k_scale_float", 1.0) + use_fused_fp8_decode = ( + self.attention_backend in self._MLA_KERNEL_BACKENDS + and getattr(ctx.attn_backend, "data_type", None) == torch.float8_e4m3fn + and self.rotary_emb is not None + and k_scale == 1.0 + ) + + if use_fused_fp8_decode: + q_nope_absorbed = Q[..., : self.kv_lora_rank] + k_nope_raw = K[..., : self.kv_lora_rank] + k_pe_raw = K[..., self.kv_lora_rank :] + + query_fp8, key_fp8 = apply_rope_mla( + positions=positions, + q_rope=q_pe, + k_rope=k_pe_raw, + q_nope=q_nope_absorbed, + k_nope=k_nope_raw, + cos_sin_cache=self.rotary_emb.cos_sin_cache, + is_neox=getattr(self.rotary_emb, "is_neox_style", True), + quant_scale_q=1.0, + quant_scale_kv=k_scale, + enable_pdl=pdl_enabled(), + ) + + # Write FP8 KV cache (single write, no double-write) + ctx.token_to_kv_pool.set_mla_kv_buffer( + self.attn_mqa, + out_cache_loc, + cache_k_nope=key_fp8[..., : self.kv_lora_rank], + cache_k_rope=key_fp8[..., self.kv_lora_rank :], + ) + return query_fp8, key_fp8 + + elif self.rotary_emb is not None and q_nope.size(0) > 0: + # Apply RoPE directly on Q and K slices + q_pe, k_pe = self.rotary_emb( + positions, + q_pe, + K[..., self.kv_lora_rank :], + fused_set_kv_buffer_arg=( + create_fused_set_kv_buffer_arg( + value=K[..., : self.kv_lora_rank], + layer=self.attn_mqa, + out_cache_loc=out_cache_loc, + token_to_kv_pool=ctx.token_to_kv_pool, + ) + if self.use_fused_set_kv_buffer + else None + ), + ) + Q[..., self.kv_lora_rank :].copy_(q_pe) + K[..., self.kv_lora_rank :].copy_(k_pe) + else: + Q[..., self.kv_lora_rank :] = q_pe + + # For MLA kernel backends, write KV cache here (model-owned) so the + # backend never has to. This unifies the FP8 fused path (written above) + # and the BF16 path into a single ownership model. + if ( + self.attention_backend in self._MLA_KERNEL_BACKENDS + and not self.use_fused_set_kv_buffer + ): + ctx.token_to_kv_pool.set_mla_kv_buffer( + self.attn_mqa, + out_cache_loc, + cache_k_nope=K[..., : self.kv_lora_rank], + cache_k_rope=K[..., self.kv_lora_rank :], + ) + + return Q, K + + def forward_absorb_attn_v_proj( + self, + Q, + K, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + output: torch.Tensor, + record_kv_cache: bool | None = None, + ) -> torch.Tensor: + # MLA kernel backends: KV cache already written in forward_absorb_qkv_proj. + # Other backends: write via fused_set_kv_buffer or let backend handle it. + if self.attention_backend in self._MLA_KERNEL_BACKENDS: + need_save_kv = False + else: + need_save_kv = not self.use_fused_set_kv_buffer + + attn_output = self.attn_mqa( + Q, + K, + K[..., : self.kv_lora_rank], + ctx, + out_cache_loc, + save_kv_cache=need_save_kv, + record_kv_cache=record_kv_cache, + ) + attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank) + if _is_amd: + projected = torch.bmm( + attn_output.transpose(0, 1).contiguous(), + self.w_vc.contiguous(), + ) + output.copy_(projected.transpose(0, 1).reshape_as(output)) + else: + output_view = output.view(-1, self.num_local_heads, self.v_head_dim) + torch.bmm( + attn_output.transpose(0, 1), + self.w_vc, + out=output_view.transpose(0, 1), + ) + return output + + def forward_normal_chunked( + self, + positions: torch.Tensor, + q: torch.Tensor, + latent_cache: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + output: torch.Tensor, + ) -> torch.Tensor: + # Prefill-graph padding contract: zero garbage rows the per-row projections + # + FP8 quantize would otherwise touch (see scrub_padding_tail). + ntok = sum(ctx.attn_backend.chunked_prefill_metadata.extend_seq_lens_cpu) + scrub_padding_tail(ntok, q, latent_cache) + q, k, v = self.forward_normal_chunked_kv_prepare( + positions, q, latent_cache, ctx, out_cache_loc + ) + return self.forward_normal_chunked_kv_core(q, k, v, ctx, out_cache_loc, output) + + def forward_normal_chunked_kv_prepare( + self, + positions: torch.Tensor, + q: torch.Tensor, + latent_cache: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + kv_a, k_pe = latent_cache.split( + [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + k_pe = k_pe.unsqueeze(1) + + q = q.view(-1, self.num_local_heads, self.qk_head_dim) + q_nope, q_pe = q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + kv = self.kv_b_proj(kv_a)[0] + kv = kv.view(-1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim) + k_nope = kv[..., : self.qk_nope_head_dim] + v = kv[..., self.qk_nope_head_dim :] + + # FP8 prefill: fused RoPE + FP8 quantize, direct FP8 KV cache write. + # Disabled when k_scale != 1.0; mla_fp8_utils.py documents the current limitation. + k_scale = getattr(self.attn_mha, "k_scale_float", 1.0) + use_fp8_prefill = ( + self.attention_backend in self._MLA_KERNEL_BACKENDS + and getattr(ctx.attn_backend, "data_type", None) == torch.float8_e4m3fn + and self.rotary_emb is not None + and k_scale == 1.0 + ) + + if use_fp8_prefill: + # Expand k_pe from [tokens,1,rope] to [tokens,heads,rope] for GQA + k_pe_expanded = k_pe.expand(-1, self.num_local_heads, -1) + + q_fp8, k_fp8 = apply_rope_mla( + positions=positions, + q_rope=q_pe, + k_rope=k_pe_expanded, + q_nope=q_nope, + k_nope=k_nope, + cos_sin_cache=self.rotary_emb.cos_sin_cache, + is_neox=getattr(self.rotary_emb, "is_neox_style", True), + quant_scale_q=1.0, + quant_scale_kv=k_scale, + enable_pdl=pdl_enabled(), + ) + + v_fp8 = fp8_quantize(v, enable_pdl=pdl_enabled()) + + # Write FP8 KV cache directly (skip BF16→FP8 conversion in pool) + k_pe_for_cache = k_fp8[:, 0:1, self.qk_nope_head_dim :] + kv_a_fp8 = fp8_quantize(kv_a, enable_pdl=pdl_enabled()) + ctx.token_to_kv_pool.set_mla_kv_buffer( + self.attn_mha, + out_cache_loc, + cache_k_nope=kv_a_fp8.unsqueeze(1), + cache_k_rope=k_pe_for_cache, + ) + + return q_fp8, k_fp8, v_fp8 + + # BF16 path: apply RoPE, assemble Q/K, write cache + if self.rotary_emb is not None: + q_pe, k_pe = self.rotary_emb( + positions, + q_pe, + k_pe, + fused_set_kv_buffer_arg=( + create_fused_set_kv_buffer_arg( + value=kv_a.unsqueeze(1), + layer=self.attn_mha, + out_cache_loc=out_cache_loc, + token_to_kv_pool=ctx.token_to_kv_pool, + ) + if self.use_fused_set_kv_buffer + else None + ), + ) + + q[..., self.qk_nope_head_dim :] = q_pe + k = torch.empty_like(q) + k[..., : self.qk_nope_head_dim] = k_nope + k[..., self.qk_nope_head_dim :] = k_pe + + if not self.use_fused_set_kv_buffer: + ctx.token_to_kv_pool.set_mla_kv_buffer( + self.attn_mha, + out_cache_loc, + cache_k_nope=kv_a.unsqueeze(1), + cache_k_rope=k_pe, + ) + + return q, k, v + + def forward_normal_chunked_kv_core( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + output: torch.Tensor, + ) -> torch.Tensor: + attn_backend = ctx.attn_backend + chunk_meta = attn_backend.chunked_prefill_metadata + token_to_kv_pool = ctx.token_to_kv_pool + + # Scale compensation for FP8 prefill: bmm1_scale = k_scale * softmax_scale + scaling = self.attn_mha.scaling + k_scale = getattr(self.attn_mha, "k_scale_float", 1.0) + if q.dtype == torch.float8_e4m3fn: + scaling = k_scale * scaling + + # Causal self-attention over the new chunk tokens. q_lens == kv_lens == + # extend_seq_lens, so cum_seq_lens_q and cum_seq_lens_kv alias the same + # cum_extend_seq_lens. Causal pass writes directly into output; each + # chunk's merge accumulates in place via attn_merge_state(inplace=True). + num_extends = chunk_meta.extend_seq_lens.size(0) + output_view = output.view(-1, self.num_local_heads, self.v_head_dim) + _, accum_lse = attn_backend.forward_extend_chunked( + q, + k, + v, + scaling, + self.attn_mha.logit_cap, + cum_seq_lens_q=chunk_meta.cum_extend_seq_lens, + cum_seq_lens_kv=chunk_meta.cum_extend_seq_lens, + max_q_len=chunk_meta.max_extend_seq_len, + max_kv_len=chunk_meta.max_extend_seq_len, + seq_lens=chunk_meta.extend_seq_lens, + batch_size=num_extends, + causal=True, + out=output_view, + ) + + # Always read KV cache as BF16 for kv_b_proj (weight is BF16), even if Q is FP8. + read_dtype = ( + q.dtype + if q.dtype not in (torch.float8_e4m3fn, torch.float8_e5m2) + else torch.bfloat16 + ) + + for loop_idx in range(chunk_meta.chunked_loop_num): + chunk_kv_indices = chunk_meta.chunk_kv_indices_list[loop_idx] + + kv_a_normed, k_pe = token_to_kv_pool.get_mla_kv_buffer( + self.attn_mha, chunk_kv_indices, read_dtype + ) + + kv_a_normed = kv_a_normed.squeeze(1) + kv = self.kv_b_proj(kv_a_normed)[0] + kv = kv.view( + -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim + ) + v = kv[..., self.qk_nope_head_dim :] + k_nope = kv[..., : self.qk_nope_head_dim] + + if q.dtype == torch.float8_e4m3fn: + # FP8 Attention + k, v = mla_kv_pack_quantize_fp8( + k_nope, k_pe, v, k_scale_inv=1.0 / k_scale, enable_pdl=pdl_enabled() + ) + else: + # BF16 Attention + k = torch.cat( + [k_nope, k_pe.expand(-1, self.num_local_heads, -1)], dim=-1 + ) + + chunk_output, lse = attn_backend.forward_extend_chunked( + q, + k, + v, + scaling, + self.attn_mha.logit_cap, + cum_seq_lens_q=chunk_meta.cum_extend_seq_lens, + cum_seq_lens_kv=chunk_meta.cu_chunked_seq_len[loop_idx], + max_q_len=chunk_meta.max_extend_seq_len, + max_kv_len=chunk_meta.max_chunk_len_per_loop[loop_idx], + seq_lens=chunk_meta.chunked_seq_len[loop_idx], + batch_size=num_extends, + causal=False, + ) + + attn_merge_state( + output_view, + accum_lse, + chunk_output, + lse, + inplace=True, + enable_pdl=pdl_enabled(), + ) + + return output + + +class DeepseekV3DraftAttentionMLA(DeepseekV3AttentionMLA): + """Draft variant of MLA shared by the NextN and Eagle3 MLA drafters. + + On the active first draft step the full ``latent_cache`` (N rows) is + projected so every KV cache entry is written, but only the live query rows + (``ctx.gather_ids``) run the absorbed decode attention, narrowing the output + to ``[bs, H]``. Multi-step decode and target paths delegate to the base. + Single-layer only, so dropping the dead rows has no downstream consumer. + """ + + def _attn( + self, + positions: torch.Tensor, + q: torch.Tensor, + latent_cache: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + ) -> torch.Tensor: + if ctx.accept_lengths is None: + return super()._attn(positions, q, latent_cache, ctx, out_cache_loc) + + self._apply_correction(ctx) + + # Full q/latent_cache write all KV cache rows; only the live rows + # (ctx.gather_ids) run the absorbed decode attention, so the output is + # narrowed to [bs, H] for o_proj / MLP / post-norms. + decode_ctx = replace(ctx, forward_mode=ForwardMode.DECODE) + Q, K = self.forward_absorb_qkv_proj( + q, + latent_cache, + positions, + decode_ctx, + out_cache_loc, + ) + Q = Q.index_select(0, ctx.gather_ids) + attn_output = q.new_empty(ctx.bs, self.num_local_heads * self.v_head_dim) + # gather_ids keeps one live row per request, so the decode runs on the + # full bs -- the page table and seq lens must span the same rows. Drop + # the [num_extends:] slice a MIXED target's first-step metadata sets up + # (mirrors the multi-step drafter loop's override_num_extends(0)). + with ctx.attn_backend.override_num_extends(0): + self.forward_absorb_attn_v_proj( + Q, + K, + decode_ctx, + out_cache_loc, + attn_output, + # Real-mode record: decode_ctx would skip the PD cache-step here. + record_kv_cache=not ctx.forward_mode.is_decode_or_idle(), + ) + return attn_output + + def _apply_correction(self, ctx: ForwardContext) -> None: + """Trim decode rows' cache_seqlens by ``spec_num_tokens - accept_lengths``.""" + seq_lens_buf = ctx.draft_seq_lens_buf + if seq_lens_buf is None or ctx.accept_lengths is None: + return + num_extends = ctx.num_extends + if num_extends >= ctx.bs: + return + correction = ( + ctx.attn_backend.spec_num_tokens - ctx.accept_lengths[num_extends:] + ).to(seq_lens_buf.dtype) + seq_lens_buf[num_extends : ctx.bs].sub_(correction).clamp_(min=1) + + +class DeepseekV3DecoderLayer(nn.Module): + @property + def attention_cls(self) -> type[nn.Module]: + return DeepseekV3AttentionMLA + + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + is_nextn: bool = False, + prefix: str = "", + alt_stream: torch.cuda.Stream | None = None, + ) -> None: + super().__init__() + self.mapping = mapping + self.hidden_size = config.hidden_size + rope_theta = get_rope_theta(config) + rope_scaling = getattr(config, "rope_scaling", None) + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + + self.self_attn = self.attention_cls( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + qk_nope_head_dim=config.qk_nope_head_dim, + qk_rope_head_dim=config.qk_rope_head_dim, + v_head_dim=config.v_head_dim, + q_lora_rank=( + config.q_lora_rank if hasattr(config, "q_lora_rank") else None + ), + kv_lora_rank=config.kv_lora_rank, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=( + None + if "self_attn" in getattr(config, "disable_quant_module", []) + else quant_config + ), + layer_id=layer_id, + prefix=add_prefix("self_attn", prefix), + reduce_attn_results=False, + alt_stream=alt_stream, + mapping=self.mapping, + ) + + self.layer_id = layer_id + self.is_moe_layer = self._is_moe_layer(layer_id, is_nextn, config) + if self.is_moe_layer: + self.mlp = DeepseekV3MoE( + config=config, + mapping=self.mapping, + quant_config=quant_config, + layer_index=layer_id, + prefix=add_prefix("mlp", prefix), + alt_stream=alt_stream, + ) + else: + self.mlp = DeepseekV3MLP( + hidden_size=config.hidden_size, + intermediate_size=( + config.ffn_hidden_size + if hasattr(config, "ffn_hidden_size") + else config.intermediate_size + ), + hidden_act=config.hidden_act, + mapping=self.mapping, + quant_config=( + None + if "dense_mlp" in getattr(config, "disable_quant_module", []) + else quant_config + ), + prefix=add_prefix("mlp", prefix), + is_shared_expert=False, + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.comm_manager = CommManager( + mapping=self.mapping, + layer_id=self.layer_id, + is_moe=self.is_moe_layer, + prev_is_moe=self._is_moe_layer(layer_id - 1, is_nextn, config), + input_layernorm=self.input_layernorm, + post_attn_layernorm=self.post_attention_layernorm, + ) + + @staticmethod + def _is_moe_layer(layer_id: int, is_nextn: bool, config): + if is_nextn: + return True + if ( + config.n_routed_experts is not None + and layer_id >= config.first_k_dense_replace + and layer_id % config.moe_layer_freq == 0 + ): + return True + return False + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + residual: torch.Tensor | None, + ) -> torch.Tensor: + + num_global_tokens, max_num_tokens_per_gpu = self.comm_manager.get_num_tokens( + ctx + ) + + if not ctx.forward_mode.is_idle(): + hidden_states, residual = self.comm_manager.input_reduce_norm( + hidden_states, residual + ) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ctx=ctx, + out_cache_loc=out_cache_loc, + comm_manager=self.comm_manager, + ) + hidden_states, residual = self.comm_manager.post_attn_reduce_norm( + hidden_states, residual, ctx + ) + hidden_states = self.forward_mlp( + hidden_states, + residual, + ctx, + num_global_tokens, + max_num_tokens_per_gpu, + ) + else: + hidden_states = self.forward_mlp( + hidden_states, + residual, + ctx, + num_global_tokens, + max_num_tokens_per_gpu, + ) + return hidden_states, residual + + def input_layer_norm_fn(self, hidden_states, residual): + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + return hidden_states, residual + + def forward_mlp( + self, + hidden_states, + residual, + ctx: ForwardContext, + num_global_tokens, + max_num_tokens_per_gpu, + ): + hidden_states = self.comm_manager.pre_mlp_comm(hidden_states, ctx) + if self.is_moe_layer: + hidden_states = self.mlp( + hidden_states, num_global_tokens, max_num_tokens_per_gpu + ) + else: + hidden_states = self.mlp(hidden_states) + hidden_states, residual = self.comm_manager.post_mlp_fused( + hidden_states, residual, ctx + ) + return hidden_states + + +class DeepseekV3Model(nn.Module): + fall_back_to_pt_during_load = False + + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.mapping = mapping + self.padding_id = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + ) + self.alt_stream = torch.cuda.Stream() + # config.num_hidden_layers = 5; self.start_layer,self.end_layer = 0, 5 + self.layers = nn.ModuleList( + [ + DeepseekV3DecoderLayer( + config, + layer_id, + mapping=self.mapping, + quant_config=quant_config, + prefix=add_prefix(f"layers.{layer_id}", prefix), + alt_stream=self.alt_stream, + ) + for layer_id in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + # For EAGLE3 support: set of layer indices whose *input* hidden states + # are captured. Populated by set_eagle3_layers_to_capture(). + self.layers_to_capture: set = set() + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + input_embeds: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, list[torch.Tensor] | None]: + if input_embeds is not None: + hidden_states = input_embeds + else: + hidden_states = self.embed_tokens(input_ids) + if CP_METADATA: + hidden_states = cp_split_and_rebuild_data( + hidden_states, + CP_METADATA.value.split_list, + CP_METADATA.value.zigzag_index, + ) + positions = cp_split_and_rebuild_data( + positions, CP_METADATA.value.split_list, CP_METADATA.value.zigzag_index + ) + residual = None + aux_hidden_states = [] if self.layers_to_capture else None + for i in range(len(self.layers)): + if aux_hidden_states is not None and i in self.layers_to_capture: + # Under RSAG the inter-layer hidden/residual are reduce- + # scattered across the attn TP group; aux consumers (e.g. the + # EAGLE3 drafter) expect full rows, so gather before capturing. + aux = ( + hidden_states + residual if residual is not None else hidden_states + ) + gathered = self.layers[i].comm_manager.gather_residual(aux, ctx) + aux_hidden_states.append( + gathered if gathered is aux else gathered.clone() + ) + with get_global_expert_distribution_recorder().with_current_layer(i): + layer = self.layers[i] + hidden_states, residual = layer( + positions, + hidden_states, + ctx, + out_cache_loc, + residual, + ) + if not ctx.forward_mode.is_idle(): + if not ENABLE_CP: + hidden_states, _ = layer.comm_manager.final_norm( + hidden_states, residual, ctx, self.norm + ) + else: + hidden_states, _ = self.norm(hidden_states, residual) + if CP_METADATA: + hidden_states = cp_all_gather_rerange_output( + hidden_states, + CP_METADATA.value, + self.mapping.attn.tp_rank, + self.mapping.attn.tp_group, + ) + return hidden_states, aux_hidden_states + + +class DeepseekV3ForCausalLM(BaseCausalLM): + model_cls = DeepseekV3Model + + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + model: DeepseekV3Model | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + self._model_override = model + super().__init__( + config=config, + mapping=mapping, + quant_config=quant_config, + prefix=prefix, + ) + + def resolve_model( + self, + config: PretrainedConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None, + prefix: str, + ) -> DeepseekV3Model: + if self._model_override is not None: + return self._model_override + return self.model_cls( + config, + mapping=mapping, + quant_config=quant_config, + prefix=add_prefix("model", prefix), + ) + + def post_init(self) -> None: + self._routed_experts_weights_of_layer = LazyValue( + lambda: { + layer_id: layer.mlp.get_moe_routed_weights() + for layer_id, layer in enumerate(self.model.layers) + if isinstance(layer.mlp, DeepseekV3MoE) + } + ) + + @property + def routed_experts_weights_of_layer(self): + return self._routed_experts_weights_of_layer.value + + def set_eagle3_layers_to_capture(self, layer_ids: list[int] | None = None): + # layer_ids are 0-indexed from the external API; +1 because the capture + # check runs *before* the layer forward, so index i captures layer i-1's output. + if layer_ids is None: + num_layers = self.config.num_hidden_layers + self.model.layers_to_capture = {2, num_layers // 2, num_layers - 3} + else: + self.model.layers_to_capture = {val + 1 for val in layer_ids} + + def set_dflash_layers_to_capture(self, layer_ids: list[int]): + # DFlash checkpoints name 0-indexed target layer outputs. The capture + # check runs before layer i, so capture at i + 1 for layer i's output. + num_layers = len(self.model.layers) + if len(set(layer_ids)) != len(layer_ids): + raise ValueError("DFLASH target_layer_ids must be unique.") + + invalid = [val for val in layer_ids if val < 0 or val + 1 >= num_layers] + if invalid: + raise ValueError( + "DFLASH target_layer_ids must map to capturable target layer " + f"outputs. Got invalid ids {invalid}; valid range is " + f"[0, {num_layers - 2}] for {num_layers} target layers." + ) + self.model.layers_to_capture = {val + 1 for val in layer_ids} + + def get_param(self, params_dict, name): + if name in params_dict: + return params_dict[name] + + if "language_model." in name: + name = name.replace("language_model.", "") + if name in params_dict: + return params_dict[name] + + if name.endswith(_OPTIONAL_MISSING_WEIGHT_SUFFIXES): + return None + + logger.warning("The %s is not in the model.", name) + return None + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + # Fuse q_a_proj and kv_a_proj_with_mqa along output dimension when q_lora_rank is not None + fuse_qkv_a_proj = getattr(self.config, "q_lora_rank", None) is not None + + params_dict = dict(self.named_parameters()) + moe_params_dict = dict(params_dict) + for param_name, param in params_dict.items(): + if param_name.startswith("model."): + moe_params_dict.setdefault( + param_name.replace("model.", "model.language_model.", 1), + param, + ) + moe_params_dict.setdefault( + param_name.replace("model.", "language_model.model.", 1), + param, + ) + # MoE expert weights, scales, and activation scales are handled + # by the checkpoint loader. + moe_loader = build_moe_checkpoint_loader( + params_dict=moe_params_dict, + expert_schema=ExpertCheckpointSchema( + gate_proj_name="gate_proj", + down_proj_name="down_proj", + up_proj_name="up_proj", + ), + num_experts=self.config.n_routed_experts, + ep_rank=self.mapping.moe.ep_rank, + ep_size=self.mapping.moe.ep_size, + ) + for name, loaded_weight in weights: + layer_id = get_layer_id(name) + if ( + layer_id is not None + and hasattr(self.model, "start_layer") + and ( + layer_id < self.model.start_layer + or layer_id >= self.model.end_layer + ) + ): + continue + if hasattr(self.config, "num_nextn_predict_layers"): + num_nextn_layers = self.config.num_nextn_predict_layers + if num_nextn_layers > 0 and name.startswith("model.layers"): + name_list = name.split(".") + if ( + len(name_list) >= 3 + and int(name_list[2]) >= self.config.num_hidden_layers + ): + continue + if "rotary_emb.inv_freq" in name: + continue + if ".indexer." in name: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + # Skip non-stacked layers and experts (experts handled below). + if weight_name not in name: + continue + # We have mlp.experts[0].gate_proj in the checkpoint. + # Since moe_loader handles the experts below, + # we need to skip here BEFORE we update the name, otherwise + # name will be updated to mlp.experts[0].gate_up_proj, which + # will then be updated below by moe_loader + # for mlp.experts[0].gate_gate_up_proj, which breaks load. + if ("mlp.experts." in name) and name not in params_dict: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = self.get_param(params_dict, name) + if param is None: + continue + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if moe_loader.matches(name): + moe_loader.load(name, loaded_weight) + continue + + if fuse_qkv_a_proj and ( + "q_a_proj" in name or "kv_a_proj_with_mqa" in name + ): + quant_block_size = 1 + # ``weight_block_size`` exists only on block-FP8 configs; + # elsewhere (e.g. compressed-tensors INT4) q/kv_a_proj is unquantized. + weight_block_size = getattr( + self.quant_config, "weight_block_size", None + ) + if weight_block_size is not None: + quant_block_size = weight_block_size[0] + begin_size_mp = { + "q_a_proj": 0, + "kv_a_proj_with_mqa": self.config.q_lora_rank, + } + if "q_a_proj" in name: + param = self.get_param( + params_dict, + name.replace("q_a_proj", "fused_qkv_a_proj_with_mqa"), + ) + weight_loader = param.weight_loader + begin_size = begin_size_mp["q_a_proj"] + elif "kv_a_proj_with_mqa" in name: + param = self.get_param( + params_dict, + name.replace( + "kv_a_proj_with_mqa", "fused_qkv_a_proj_with_mqa" + ), + ) + weight_loader = param.weight_loader + begin_size = begin_size_mp["kv_a_proj_with_mqa"] + if "scale_inv" in name: + begin_size //= quant_block_size + weight_loader(param, loaded_weight, begin_size=begin_size) + else: + # Owned-expert weights were already consumed by ``moe_loader.load(...)`` above (matches() == True branch). + # Anything reaching here that still looks like an expert weight is for an expert this rank does ot own under ep_size > 1. + if ".mlp.experts." in name: + continue + if "q_a_proj" in name and name not in params_dict: + name = name.replace("q_a_proj", "q_proj") + param = self.get_param(params_dict, name) + if param is None: + continue + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + + self.post_load_weights() + + def post_load_weights(self): + for layer_id in range(self.config.num_hidden_layers): + self_attn = self.model.layers[layer_id].self_attn + if hasattr( + self.quant_config, "weight_block_size" + ) and self_attn.kv_b_proj.weight.dtype in ( + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + ): + weight_block_size = self.quant_config.weight_block_size + if weight_block_size is not None: + if not hasattr(self_attn.kv_b_proj, "weight_scale_inv"): + raise RuntimeError( + "kv_b_proj.weight_scale_inv is required for block FP8 dequant." + ) + dtype = torch.get_default_dtype() + w = block_dequant( + self_attn.kv_b_proj.weight, + self_attn.kv_b_proj.weight_scale_inv, + weight_block_size, + ).to(dtype) + else: + w = self_attn.kv_b_proj.weight + + self_attn.w_kc, self_attn.w_vc = _prepare_mla_kv_b_proj_weights( + w, self_attn + ) + + def load_kv_cache_scales(self, quantization_param_path: str) -> None: + tp_size = self.mapping.attn.tp_size + tp_rank = self.mapping.attn.tp_rank + for layer_idx, scaling_factor in kv_cache_scales_loader( + quantization_param_path, + tp_rank, + tp_size, + self.config.num_hidden_layers, + self.config.__class__.model_type, + ): + if not isinstance(self.model.layers[layer_idx], nn.Identity): + self_attn = self.model.layers[layer_idx].self_attn + # Set on both attn_mha (non-absorbed prefill) and attn_mqa (absorbed decode). + for attn in (self_attn.attn_mha, self_attn.attn_mqa): + if attn is not None and hasattr(attn, "k_scale"): + attn.k_scale = scaling_factor + attn.k_scale_float = scaling_factor + + def get_embed_and_head(self): + return self.model.embed_tokens.weight, self.lm_head.weight + + def set_embed_and_head(self, embed, head): + del self.model.embed_tokens.weight + del self.lm_head.weight + self.model.embed_tokens.weight = embed + self.lm_head.weight = head + torch.cuda.empty_cache() + torch.cuda.synchronize() + + @classmethod + def get_model_config_for_expert_location(cls, config): + return ModelConfigForExpertLocation( + num_layers=config.num_hidden_layers, + num_logical_experts=config.n_routed_experts, + num_groups=config.n_group, + ) + + +# --------------------------------------------------------------------------- +# Eagle3 MLA draft model +# --------------------------------------------------------------------------- + + +class Eagle3MlaDecoderLayer(nn.Module): + """Single decoder layer for Eagle3 MLA draft model. + + The fused_qkv_a_proj_with_mqa is overridden to accept 2x hidden_size + input (concatenated [embeds, hidden_states]) while keeping o_proj at + the standard hidden_size output. + """ + + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + layer_id: int = 0, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.mapping = mapping + self.hidden_size = config.hidden_size + self.layer_id = layer_id + rope_theta = get_rope_theta(config) + rope_scaling = getattr(config, "rope_scaling", None) + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + + self.self_attn = DeepseekV3DraftAttentionMLA( + config=config, + mapping=self.mapping, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + qk_nope_head_dim=getattr(config, "qk_nope_head_dim", 128), + qk_rope_head_dim=getattr(config, "qk_rope_head_dim", 64), + v_head_dim=getattr(config, "v_head_dim", 128), + q_lora_rank=getattr(config, "q_lora_rank", None), + kv_lora_rank=getattr(config, "kv_lora_rank", 512), + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + layer_id=layer_id, + prefix=add_prefix("self_attn", prefix), + reduce_attn_results=False, + ) + + if hasattr(self.self_attn, "fused_qkv_a_proj_with_mqa"): + q_lora_rank = getattr(config, "q_lora_rank", 0) or 0 + kv_lora_rank = getattr(config, "kv_lora_rank", 512) + qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 64) + self.self_attn.fused_qkv_a_proj_with_mqa = DeepseekV3FusedQkvAProjWithMqa( + 2 * self.hidden_size, + q_lora_rank + kv_lora_rank + qk_rope_head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix( + "fused_qkv_a_proj_with_mqa", + add_prefix("self_attn", prefix), + ), + ) + + self.mlp = DeepseekV3MLP( + hidden_size=config.hidden_size, + intermediate_size=getattr( + config, "intermediate_size", config.hidden_size * 4 + ), + hidden_act=getattr(config, "hidden_act", "silu"), + mapping=self.mapping, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + + self.hidden_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.fused_input_hidden_norm = FusedRMSNorm( + self.input_layernorm, + self.hidden_norm, + ) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + self.comm_manager = CommManager( + mapping=self.mapping, + layer_id=self.layer_id, + is_moe=False, + prev_is_moe=False, + post_attn_layernorm=self.post_attention_layernorm, + ) + + def forward( + self, + positions: torch.Tensor, + embeds: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + residual = hidden_states + + if not ctx.forward_mode.is_idle(): + fused_norm_out = torch.empty( + embeds.size(0), + self.hidden_size * 2, + dtype=embeds.dtype, + device=embeds.device, + ) + # FusedRMSNorm's q_a/kv_a kwargs are MLA-specific names. + # Here embeds and hidden_states corresponds to q_a and kv_a, separately. + self.fused_input_hidden_norm( + input_q_a=embeds, + input_kv_a=hidden_states, + output_q_a=fused_norm_out[..., : self.hidden_size], + output_kv_a=fused_norm_out[..., self.hidden_size :], + ) + + hidden_states = self.self_attn( + positions=positions, + hidden_states=fused_norm_out, + ctx=ctx, + out_cache_loc=out_cache_loc, + comm_manager=self.comm_manager, + ) + + # Active first draft step narrows attn output to [bs, H]; align the + # residual to the same live rows before the post-attn reduce-norm. + if ctx.accept_lengths is not None: + residual = residual.index_select(0, ctx.gather_ids) + hidden_states, residual = self.comm_manager.post_attn_reduce_norm( + hidden_states, residual, ctx + ) + + hidden_states = self.comm_manager.pre_mlp_comm(hidden_states, ctx) + hidden_states = self.mlp(hidden_states) + hidden_states, residual = self.comm_manager.post_mlp_fused( + hidden_states, residual, ctx + ) + + return hidden_states, residual + + +class Eagle3MlaModel(nn.Module): + @staticmethod + def _get_eagle_layer_ids(config: PretrainedConfig): + """Extract eagle aux hidden state layer IDs from config, or None if absent.""" + eagle_config = getattr(config, "eagle_config", None) + if eagle_config is None: + return getattr(config, "eagle_aux_hidden_state_layer_ids", None) + if isinstance(eagle_config, dict): + return eagle_config.get("eagle_aux_hidden_state_layer_ids", None) + return getattr(eagle_config, "eagle_aux_hidden_state_layer_ids", None) + + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.mapping = mapping + self.config = config + self.vocab_size = config.vocab_size + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + prefix=add_prefix("embed_tokens", prefix), + ) + + layer_ids = self._get_eagle_layer_ids(config) + self.num_fc_input_dim = len(layer_ids) if layer_ids is not None else 3 + + target_hidden_size = getattr(config, "target_hidden_size", config.hidden_size) + fc_input_size = target_hidden_size * self.num_fc_input_dim + + self.fc = ColumnParallelLinear( + fc_input_size, + config.hidden_size, + bias=False, + gather_output=True, + quant_config=quant_config, + prefix=add_prefix("fc", prefix), + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + + self.midlayer = Eagle3MlaDecoderLayer( + config, + mapping=self.mapping, + layer_id=0, + quant_config=quant_config, + prefix=prefix, + ) + + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + input_embeds: torch.Tensor | None = None, + captured_hidden_states: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, list[torch.Tensor]]: + if captured_hidden_states is None: + raise ValueError("Eagle3 MLA forward requires captured_hidden_states.") + if input_embeds is None: + embeds = self.embed_tokens(input_ids) + else: + embeds = input_embeds + + hidden_states = captured_hidden_states + if hidden_states.size(-1) != embeds.size(-1): + hidden_states, _ = self.fc(hidden_states) + + residual = None + hidden_states, residual = self.midlayer( + positions, + embeds, + hidden_states, + ctx, + out_cache_loc, + residual, + ) + + comm_manager = self.midlayer.comm_manager + if comm_manager.should_fuse(hidden_states.size(0)): + hidden_states_to_logits, hidden_states_to_aux, *_ = ( + self.norm.forward_with_allreduce_fusion( + self.mapping.dense.tp_rank, + self.mapping.dense.tp_group, + hidden_states, + residual, + ) + ) + else: + hidden_states_to_logits, hidden_states_to_aux = self.norm( + hidden_states, residual + ) + hidden_states_to_logits, _ = comm_manager.post_final_norm_comm( + hidden_states_to_logits, None, ctx + ) + hidden_states_to_aux, _ = comm_manager.post_final_norm_comm( + hidden_states_to_aux, None, ctx + ) + return hidden_states_to_logits, [hidden_states_to_aux] + + +class Eagle3DeepseekV2ForCausalLM(DeepseekV3ForCausalLM): + """Eagle3 MLA draft model for DeepSeek-V2/V3 / Kimi-K2 style architectures. + + Inherits weight-loading fusion logic from DeepseekV3ForCausalLM but uses + Eagle3MlaModel internally with a single MLA decoder layer that accepts + concatenated [embeds || hidden_states] as input. + """ + + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + nn.Module.__init__(self) + self.config = config + self.mapping = mapping + self.quant_config = quant_config + + if self.config.num_hidden_layers != 1: + raise ValueError("Eagle3 MLA drafter currently only supports 1 layer") + + self.model = Eagle3MlaModel( + config, + mapping=self.mapping, + quant_config=quant_config, + prefix=add_prefix("model", prefix), + ) + + self.load_lm_head_from_target = False + if self.config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + draft_vocab_size = ( + getattr(config, "draft_vocab_size", None) or config.vocab_size + ) + if not hasattr(config, "draft_vocab_size"): + self.load_lm_head_from_target = True + self.lm_head = ParallelLMHead( + draft_vocab_size, + config.hidden_size, + quant_config=quant_config, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + prefix=add_prefix("lm_head", prefix), + ) + + self.logits_processor = LogitsProcessor( + config, + skip_all_gather=self.mapping.attn.has_dp, + do_argmax=True, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + self.capture_aux_hidden_states = True + self.hot_token_id = None + + def forward( + self, + ctx: ForwardContext, + input_ids: torch.Tensor, + positions: torch.Tensor, + out_cache_loc: torch.Tensor, + **kwargs, + ) -> torch.Tensor: + with report_collective_sizing(ctx, ctx.bs, ctx.global_bs): + return super().forward(ctx, input_ids, positions, out_cache_loc, **kwargs) + + def prepare_model_kwargs( + self, ctx: ForwardContext, input_ids: torch.Tensor, kwargs: dict + ) -> dict: + model_kwargs = super().prepare_model_kwargs(ctx, input_ids, kwargs) + captured_hidden_states = kwargs.get("captured_hidden_states") + if captured_hidden_states is not None: + model_kwargs["captured_hidden_states"] = captured_hidden_states + else: + # During CUDA graph capture warmup, provide dummy hidden states. + target_hidden_size = getattr( + self.config, "target_hidden_size", self.config.hidden_size + ) + num_fc = self.model.num_fc_input_dim + model_kwargs["captured_hidden_states"] = torch.zeros( + input_ids.size(0), + target_hidden_size * num_fc, + dtype=torch.bfloat16, + device=input_ids.device, + ) + return model_kwargs + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + remapped = [] + for name, loaded_weight in weights: + if "d2t" in name: + self.hot_token_id = loaded_weight + torch.arange(loaded_weight.size(0)) + continue + if "t2d" in name: + continue + + new_name = re.sub(r"^layers\.0\.", "midlayer.", name) + + if "lm_head" not in new_name: + new_name = f"model.{new_name}" + else: + self.load_lm_head_from_target = False + remapped.append((new_name, loaded_weight)) + + super().load_weights(remapped) + + def post_load_weights(self): + self_attn = self.model.midlayer.self_attn + if ( + self.quant_config is not None + and hasattr(self.quant_config, "weight_block_size") + and self_attn.kv_b_proj.weight.dtype + in (torch.float8_e4m3fn, torch.float8_e4m3fnuz) + ): + weight_block_size = self.quant_config.weight_block_size + if weight_block_size is not None: + if not hasattr(self_attn.kv_b_proj, "weight_scale_inv"): + raise RuntimeError( + "kv_b_proj.weight_scale_inv is required for block FP8 dequant." + ) + dtype = torch.get_default_dtype() + w = block_dequant( + self_attn.kv_b_proj.weight, + self_attn.kv_b_proj.weight_scale_inv, + weight_block_size, + ).to(dtype) + else: + w = self_attn.kv_b_proj.weight + else: + w = self_attn.kv_b_proj.weight + + self_attn.w_kc, self_attn.w_vc = _prepare_mla_kv_b_proj_weights(w, self_attn) + + def get_hot_token_id(self): + return self.hot_token_id + + def set_embed_and_head(self, embed, head): + if ( + hasattr(self.config, "target_hidden_size") + and self.config.target_hidden_size != self.config.hidden_size + ): + return + del self.model.embed_tokens.weight + self.model.embed_tokens.weight = embed + if head is not None and self.load_lm_head_from_target: + del self.lm_head.weight + self.lm_head.weight = head + torch.cuda.empty_cache() + torch.cuda.synchronize() + + +EntryClass = [ + DeepseekV3ForCausalLM, + Eagle3DeepseekV2ForCausalLM, +] diff --git a/python/tokenspeed/runtime/models/deepseek_v4.py b/python/tokenspeed/runtime/models/deepseek_v4.py new file mode 100644 index 0000000..18b09b5 --- /dev/null +++ b/python/tokenspeed/runtime/models/deepseek_v4.py @@ -0,0 +1,4306 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Inference-only DeepSeek V4 model skeleton. + +This module intentionally registers only architecture pieces that map to the +DeepSeek V4 Flash checkpoint. The sparse MLA forward path still fails loudly +until the HCA/CSA cache kernels are wired into TokenSpeed. +""" + +from __future__ import annotations + +import os +import re +from collections.abc import Iterable +from dataclasses import dataclass + +import torch +import torch.nn.functional as F + +try: + # Optional dependency; the module-level wrapper imports the external + # `deep_gemm` package unguarded, which is not installed in baseline V4 + # builds. Callsites guard usage with `deep_gemm is not None`. + from tokenspeed_kernel.thirdparty import deep_gemm +except ImportError: + deep_gemm = None # type: ignore[assignment] + +from tokenspeed_kernel.ops.attention.cuda.deepseek_v4 import ( + has_indexer_mxfp4_paged_gather, + has_persistent_topk, + indexer_mxfp4_paged_gather, + persistent_topk, +) +from tokenspeed_kernel.ops.attention.triton.deepseek_v4 import ( + deepseek_v4_indexer_decode_metadata_compute, +) +from tokenspeed_kernel.platform import current_platform +from tokenspeed_kernel.thirdparty.cuda import ( + dsv3_router_gemm, + hash_softplus_sqrt_topk_flash, + softplus_sqrt_topk_flash, +) +from tokenspeed_kernel.thirdparty.triton import ( + stage_deepseek_v4_mega_moe_inputs as _stage_deepseek_v4_mega_moe_inputs, +) +from tokenspeed_kernel.thirdparty.trtllm import ( + fast_topk_v2, +) +from torch import nn +from transformers import PretrainedConfig + +from tokenspeed.runtime.configs.deepseek_v4_cache_spec import ( + DEEPSEEK_V4_MXFP4_BLOCK_SIZE, + V4_KERNEL_BLOCK_ROWS, + deepseek_v4_indexer_mxfp4_layout_from_row_bytes, + deepseek_v4_indexer_mxfp4_scale_dim, + deepseek_v4_indexer_mxfp4_value_bytes, + deepseek_v4_nope_dim, + v4_compressed_kv_group_id, +) +from tokenspeed.runtime.distributed import Mapping +from tokenspeed.runtime.distributed.comm_manager import CommManager +from tokenspeed.runtime.distributed.process_group_manager import ( + process_group_manager as pg_manager, +) +from tokenspeed.runtime.execution.breakable_cuda_graph import ( + break_point, + current_forward_ctx, + slice_to_real_tokens, +) +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.execution.cuda_graph_wrapper import get_is_capture_mode +from tokenspeed.runtime.execution.forward_batch_info import ForwardMode +from tokenspeed.runtime.layers.attention.deepseek_v4.metadata import ( + DeepseekV4ForwardMetadata, + DeepseekV4IndexerBatchMetadata, + DeepseekV4IndexerDecodePlan, + DeepseekV4IndexerPrefillChunkPlan, + DeepseekV4IndexerPrefillMetadata, + DeepseekV4SparseIndexerMetadata, +) +from tokenspeed.runtime.layers.attention.deepseek_v4_ops import ( + deepseek_v4_csa_compress_kv_cache_insert, + deepseek_v4_csa_indexer_cache_insert, + deepseek_v4_fused_inv_rope_fp8_quant, + deepseek_v4_hca_compress_kv_cache_insert, + deepseek_v4_prepare_indexer_q_mxfp4, + fused_qnorm_rope_kv_insert, + save_deepseek_v4_compressor_state, +) +from tokenspeed.runtime.layers.attention.kv_cache.deepseek_v4 import ( + _group_slot_mapping_from_raw, + _mask_invalid_graph_tokens, +) +from tokenspeed.runtime.layers.deepseek_v4_mhc import mhc_fused_hc as fast_mhc_fused_hc +from tokenspeed.runtime.layers.deepseek_v4_mhc import mhc_post as fast_mhc_post +from tokenspeed.runtime.layers.deepseek_v4_mhc import mhc_pre as fast_mhc_pre +from tokenspeed.runtime.layers.layernorm import FusedRMSNorm, RMSNorm +from tokenspeed.runtime.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from tokenspeed.runtime.layers.moe import ( + ExpertCheckpointSchema, + build_moe_checkpoint_loader, +) +from tokenspeed.runtime.layers.moe.expert import MoELayer +from tokenspeed.runtime.layers.moe.topk import ( + BypassedTopKOutput, + StandardTopKOutput, + TopK, +) +from tokenspeed.runtime.layers.moe.utils import RoutingMethodType, get_moe_backend +from tokenspeed.runtime.layers.quantization import Mxfp4Config +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.layers.rotary_embedding import get_rope +from tokenspeed.runtime.layers.vocab_parallel_embedding import VocabParallelEmbedding +from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader +from tokenspeed.runtime.models.base import BaseCausalLM +from tokenspeed.runtime.moe.expert_location import ModelConfigForExpertLocation +from tokenspeed.runtime.utils import ( + add_prefix, + get_colorful_logger, + set_weight_attrs, +) +from tokenspeed.runtime.utils.cuda_stream import StreamFork +from tokenspeed.runtime.utils.custom_ops import direct_register_custom_op +from tokenspeed.runtime.utils.env import global_server_args_dict, pdl_enabled +from tokenspeed.runtime.utils.nvtx import nvtx_range + +_platform = current_platform() + + +logger = get_colorful_logger(__name__) + + +def _deepseek_v4_metadata_matches_tokens(metadata, num_tokens: int) -> bool: + return ( + metadata is not None + and getattr(metadata, "token_to_req_indices", None) is not None + and metadata.token_to_req_indices.numel() == num_tokens + ) + + +def _deepseek_v4_indexer_token_split( + forward_mode: ForwardMode | None, + metadata, + total_tokens: int, +) -> tuple[int, int]: + if forward_mode is not None and forward_mode.is_mixed(): + return int(metadata.num_prefill_tokens), metadata.decode_token_count() + if forward_mode is not None and forward_mode.is_decode(): + return 0, int(total_tokens) + return int(total_tokens), 0 + + +def _deepseek_v4_forward_metadata(ctx: ForwardContext): + metadata = getattr(ctx.attn_backend, "forward_metadata", None) + forward_mode = getattr(ctx, "forward_mode", None) + if forward_mode is not None and forward_mode.is_extend_or_mixed(): + return getattr(ctx.attn_backend, "forward_prefill_metadata", None) or metadata + if forward_mode is not None and forward_mode.is_decode_or_idle(): + input_num_tokens = getattr(ctx, "input_num_tokens", None) + decode_metadata = getattr(ctx.attn_backend, "forward_decode_metadata", None) + if input_num_tokens is not None and _deepseek_v4_metadata_matches_tokens( + decode_metadata, + input_num_tokens, + ): + return decode_metadata + prefill_metadata = getattr(ctx.attn_backend, "forward_prefill_metadata", None) + if input_num_tokens is not None and _deepseek_v4_metadata_matches_tokens( + prefill_metadata, + input_num_tokens, + ): + return prefill_metadata + return decode_metadata or metadata or prefill_metadata + return metadata + + +def _dequant_fp8_weight(layer: nn.Module, shape: tuple[int, ...]) -> torch.Tensor: + weight = layer.weight.view(*shape) + scale = getattr(layer, "weight_scale_inv", None) + if scale is None or weight.dtype != torch.float8_e4m3fn: + return weight.float() + + block_n, block_k = getattr(layer.quant_config, "weight_block_size", (128, 128)) + if len(shape) == 2: + out_dim, in_dim = shape + scale = scale.view( + (out_dim + block_n - 1) // block_n, + (in_dim + block_k - 1) // block_k, + ) + expanded_scale = ( + scale.float() + .repeat_interleave(block_n, dim=0) + .repeat_interleave(block_k, dim=1) + ) + return weight.float() * expanded_scale[:out_dim, :in_dim] + + groups, out_dim, in_dim = shape + scale = scale.view( + groups, + (out_dim + block_n - 1) // block_n, + (in_dim + block_k - 1) // block_k, + ) + expanded_scale = ( + scale.float() + .repeat_interleave(block_n, dim=1) + .repeat_interleave(block_k, dim=2) + ) + return weight.float() * expanded_scale[:, :out_dim, :in_dim] + + +def _deepseek_v4_router_gemm( + hidden_states: torch.Tensor, + weight: torch.Tensor, +) -> torch.Tensor: + if ( + hidden_states.dim() == 2 + and hidden_states.shape[0] > 0 + and hidden_states.is_cuda + and hidden_states.dtype == torch.bfloat16 + and weight.dtype in (torch.bfloat16, torch.float32) + and (_platform.is_hopper or _platform.is_blackwell) + ): + return dsv3_router_gemm( + hidden_states, + weight, + out_dtype=torch.float32, + enable_pdl=pdl_enabled(), + ) + + x = ( + hidden_states + if hidden_states.dtype == weight.dtype + else hidden_states.to(weight.dtype) + ) + return F.linear(x, weight, None).to(torch.float32) + + +def _deepseek_v4_bf16_linear_fp32( + hidden_states: torch.Tensor, + weight: torch.Tensor, +) -> torch.Tensor | None: + if ( + hidden_states.dim() == 2 + and hidden_states.shape[0] > 0 + and hidden_states.is_cuda + and hidden_states.dtype == torch.bfloat16 + and weight.is_cuda + and weight.dtype == torch.bfloat16 + and weight.dim() == 2 + and hidden_states.shape[1] == weight.shape[1] + and (_platform.is_hopper or _platform.is_blackwell) + ): + return dsv3_router_gemm( + hidden_states, + weight, + out_dtype=torch.float32, + enable_pdl=False, + ) + return None + + +def _deepseek_v4_fused_select_experts( + router_logits: torch.Tensor, + top_k: int, + renormalize: bool, + *, + correction_bias: torch.Tensor | None = None, + hash_indices_table: torch.Tensor | None = None, + input_ids: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor] | None: + if ( + not router_logits.is_cuda + or router_logits.dim() != 2 + or top_k <= 0 + or top_k > 32 + or router_logits.dtype not in (torch.float32, torch.float16, torch.bfloat16) + ): + return None + + num_experts = router_logits.shape[1] + + topk_weights = torch.empty( + router_logits.shape[0], + top_k, + dtype=torch.float32, + device=router_logits.device, + ) + topk_ids = torch.empty( + router_logits.shape[0], + top_k, + dtype=torch.int32, + device=router_logits.device, + ) + + if num_experts not in (256, 384) or top_k != 6 or not renormalize: + return None + + logits_f32 = router_logits.float().contiguous() + + try: + if hash_indices_table is not None: + if input_ids is None: + raise ValueError("hash-routed DeepSeek V4 MoE requires input_ids") + hash_softplus_sqrt_topk_flash( + logits_f32, + input_ids.reshape(-1).to(device=router_logits.device).contiguous(), + hash_indices_table.to( + device=router_logits.device, dtype=torch.int32 + ).contiguous(), + topk_ids, + topk_weights, + 1.0, + renormalize, + ) + elif correction_bias is not None: + softplus_sqrt_topk_flash( + logits_f32, + correction_bias.to( + device=router_logits.device, dtype=torch.float32 + ).contiguous(), + topk_ids, + topk_weights, + 1.0, + renormalize, + ) + else: + return None + except (AttributeError, RuntimeError): + return None + + return topk_weights, topk_ids + + +def _deepseek_v4_reorder_c4_ape_2604(ape: torch.Tensor) -> torch.Tensor: + """Convert C4 overlap APE from checkpoint layout to runtime window layout.""" + + if ape.dim() != 2 or ape.shape[0] != 4 or ape.shape[1] % 2 != 0: + raise ValueError(f"expected C4 APE [4, even], got {tuple(ape.shape)}") + older, newer = ape.chunk(2, dim=-1) + return torch.cat([older, newer], dim=0).reshape_as(ape) + + +def mhc_pre( + residual: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_eps: float, + sinkhorn_iters: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return fast_mhc_pre( + residual, + fn, + hc_scale, + hc_base, + rms_eps, + hc_eps, + sinkhorn_iters, + ) + + +def mhc_post( + hidden_states: torch.Tensor, + residual: torch.Tensor, + post: torch.Tensor, + comb: torch.Tensor, +) -> torch.Tensor: + return fast_mhc_post(hidden_states, residual, post, comb) + + +def mhc_fused_hc( + x_prev: torch.Tensor, + residual_prev: torch.Tensor, + post_prev: torch.Tensor, + comb_prev: torch.Tensor, + fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_eps: float, + hc_eps: float, + sinkhorn_iters: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + return fast_mhc_fused_hc( + x_prev, + residual_prev, + post_prev, + comb_prev, + fn, + hc_scale, + hc_base, + rms_eps, + hc_eps, + sinkhorn_iters, + ) + + +def hc_head( + hidden_states: torch.Tensor, + hc_fn: torch.Tensor, + hc_scale: torch.Tensor, + hc_base: torch.Tensor, + rms_norm_eps: float, + hc_eps: float, +) -> torch.Tensor: + shape, dtype = hidden_states.size(), hidden_states.dtype + x = hidden_states.flatten(1).float() + rsqrt = torch.rsqrt(x.square().mean(-1, keepdim=True) + rms_norm_eps) + mixes = F.linear(x, hc_fn.float()) * rsqrt + pre = torch.sigmoid(mixes * hc_scale.float() + hc_base.float()) + hc_eps + y = torch.sum(pre.unsqueeze(-1) * x.view(shape), dim=1) + return y.to(dtype) + + +def deepseek_v4_select_experts( + router_logits: torch.Tensor, + top_k: int, + renormalize: bool, + *, + correction_bias: torch.Tensor | None = None, + hash_indices_table: torch.Tensor | None = None, + input_ids: torch.Tensor | None = None, + need_scores: bool = True, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """DeepSeek V4 MoE routing. + + DeepSeek V4 uses sqrt(softplus(logits)) as expert scores. Correction bias + only affects expert selection; the gathered expert weights come from the + unbiased scores. Hash-routed layers use checkpoint-provided expert ids but + still gather weights from the gate scores. + + Set ``need_scores=False`` when the caller discards the third return value + (e.g. mega_moe) to skip the redundant sqrt(softplus(logits)) computation. + """ + + fused_topk = _deepseek_v4_fused_select_experts( + router_logits, + top_k, + renormalize, + correction_bias=correction_bias, + hash_indices_table=hash_indices_table, + input_ids=input_ids, + ) + if fused_topk is not None: + topk_weights, topk_ids = fused_topk + if need_scores: + scores = torch.sqrt(F.softplus(router_logits.float())) + else: + scores = router_logits + return topk_weights, topk_ids, scores + + scores = torch.sqrt(F.softplus(router_logits.float())) + if hash_indices_table is not None: + if input_ids is None: + raise ValueError("hash-routed DeepSeek V4 MoE requires input_ids") + topk_ids = hash_indices_table[input_ids.reshape(-1)].to( + device=scores.device, + dtype=torch.long, + ) + else: + scores_for_choice = scores + if correction_bias is not None: + scores_for_choice = scores_for_choice + correction_bias.to( + device=scores.device, + dtype=scores.dtype, + ).unsqueeze(0) + topk_ids = torch.topk(scores_for_choice, k=top_k, dim=-1, sorted=True)[1] + + topk_weights = scores.gather(1, topk_ids) + if renormalize: + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True).clamp_min( + torch.finfo(topk_weights.dtype).tiny + ) + return topk_weights.to(torch.float32), topk_ids.to(torch.int32), scores + + +def pack_topk_as_router_logits( + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + num_experts: int, +) -> torch.Tensor: + """Encode preselected top-k weights for BYPASSED TokenSpeed MoE backends. + + MXFP4 backends currently build routing data from logits internally. Packing + the normalized top-k weights as log-probabilities with very negative values + elsewhere makes their TopK -> Softmax/Renormalize route recover the same + selected ids and weights without changing the shared backend. + """ + + router_logits = torch.full( + (topk_ids.shape[0], num_experts), + -1e20, + dtype=torch.float32, + device=topk_weights.device, + ) + safe_weights = topk_weights.clamp_min(torch.finfo(torch.float32).tiny) + router_logits.scatter_(1, topk_ids.long(), safe_weights.log()) + return router_logits + + +def _deepseek_v4_deepgemm_fp4_indexer_available(index_q: torch.Tensor) -> bool: + return ( + deep_gemm is not None + and index_q.is_cuda + and index_q.dim() >= 3 + and index_q.shape[-2] in (32, 64) + and (index_q.shape[-1] * 2) % DEEPSEEK_V4_MXFP4_BLOCK_SIZE == 0 + ) + + +def _deepseek_v4_indexer_mxfp4_cache_view( + cache_2d: torch.Tensor, + block_size: int, +) -> torch.Tensor: + row_bytes = cache_2d.shape[1] // block_size + return torch.as_strided( + cache_2d, + (cache_2d.shape[0], block_size, 1, row_bytes), + (cache_2d.stride(0), row_bytes, row_bytes, 1), + ) + + +def _deepseek_v4_gather_paged_indexer_mxfp4_cache( + cache_2d: torch.Tensor, + block_table: torch.Tensor, + cu_seq_lens: torch.Tensor, + block_size: int, + out: tuple[torch.Tensor, torch.Tensor] | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + _, value_bytes, scale_bytes = deepseek_v4_indexer_mxfp4_layout_from_row_bytes( + cache_2d.shape[1] // block_size + ) + if out is None: + total_rows = int(cu_seq_lens[-1].item()) if cu_seq_lens.numel() else 0 + values = torch.empty( + (total_rows, value_bytes), + dtype=torch.uint8, + device=cache_2d.device, + ) + scales = torch.empty( + (total_rows, scale_bytes), + dtype=torch.uint8, + device=cache_2d.device, + ) + else: + if out[0].shape[0] != out[1].shape[0]: + raise ValueError( + "DeepSeek V4 paged gather workspace value/scale rows must match, " + f"got values={out[0].shape[0]}, scales={out[1].shape[0]}" + ) + total_rows = int(out[0].shape[0]) + values = out[0][:total_rows] + scales = out[1][:total_rows] + if total_rows == 0: + return values.view(torch.int8), scales.view(torch.int32).squeeze(-1) + + if not (cache_2d.is_cuda and block_table.is_cuda and cu_seq_lens.is_cuda): + raise RuntimeError( + "DeepSeek V4 paged MXFP4 gather requires cache, block table, " + "and sequence lengths on CUDA" + ) + if not has_indexer_mxfp4_paged_gather(): + raise RuntimeError( + "DeepSeek V4 paged MXFP4 gather requires the CUDA paged gather op" + ) + indexer_mxfp4_paged_gather( + kv_cache=cache_2d, + values_out=values, + scales_out=scales, + block_table=block_table, + cu_seq_lens=cu_seq_lens, + cache_block_size=block_size, + ) + return values.view(torch.int8), scales.view(torch.int32).squeeze(-1) + + +def _deepseek_v4_indexer_topk_from_logits( + logits: torch.Tensor, + lengths: torch.Tensor, + topk_tokens: int, + *, + next_n: int = 1, + use_prefill_topk_op: bool = False, + row_starts: torch.Tensor | None = None, + row_ends: torch.Tensor | None = None, + out: torch.Tensor | None = None, + persistent_topk_workspace: torch.Tensor | None = None, +) -> torch.Tensor: + if not logits.is_cuda or logits.dtype != torch.float32: + raise RuntimeError("DeepSeek V4 indexer top-k requires CUDA float32 logits") + if not lengths.is_cuda or lengths.device != logits.device: + raise RuntimeError( + "DeepSeek V4 indexer top-k requires CUDA length tensors " + "on the logits device" + ) + if row_starts is not None and ( + not row_starts.is_cuda or row_starts.device != logits.device + ): + raise RuntimeError( + "DeepSeek V4 indexer top-k requires row_starts on the logits device" + ) + if row_ends is not None and ( + not row_ends.is_cuda or row_ends.device != logits.device + ): + raise RuntimeError( + "DeepSeek V4 indexer top-k requires row_ends on the logits device" + ) + + lengths_for_kernel = lengths.to(torch.int32).contiguous() + length_rows = lengths_for_kernel.reshape(-1) + num_tokens = length_rows.numel() + if out is None: + topk = torch.empty( + (num_tokens, topk_tokens), + device=logits.device, + dtype=torch.int32, + ) + else: + topk = out[:num_tokens] + topk.fill_(-1) + if num_tokens == 0: + return topk + max_len = logits.shape[1] if logits.dim() == 2 else 0 + if max_len <= 0: + return topk + + row_starts_for_kernel: torch.Tensor | None = None + row_ends_for_kernel: torch.Tensor | None = None + if row_starts is not None or row_ends is not None: + if row_starts is None: + row_starts_for_kernel = torch.zeros_like(length_rows) + else: + row_starts_for_kernel = row_starts.to(torch.int32).reshape(-1) + if row_ends is None: + row_ends_for_kernel = row_starts_for_kernel + length_rows + else: + row_ends_for_kernel = row_ends.to(torch.int32).reshape(-1) + length_rows = (row_ends_for_kernel - row_starts_for_kernel).clamp_min(0) + + if use_prefill_topk_op: + return _deepseek_v4_indexer_topk_from_logits_prefill_op( + logits, + length_rows, + topk_tokens, + row_starts=row_starts_for_kernel, + row_ends=row_ends_for_kernel, + out=topk, + ) + + if row_starts_for_kernel is not None or row_ends_for_kernel is not None: + raise RuntimeError( + "DeepSeek V4 decode indexer top-k does not support row ranges" + ) + if topk_tokens not in (512, 1024, 2048): + raise RuntimeError( + "DeepSeek V4 decode indexer top-k supports topk_tokens in " + "{512, 1024, 2048}" + ) + + if ( + persistent_topk_workspace is not None + and persistent_topk_workspace.is_cuda + and persistent_topk_workspace.device == logits.device + and persistent_topk_workspace.numel() >= 1024 * 1024 + and persistent_topk_workspace.dtype == torch.uint8 + ): + if not has_persistent_topk(): + raise RuntimeError( + "DeepSeek V4 persistent top-k workspace was provided, " + "but the persistent top-k kernel is unavailable" + ) + persistent_topk( + logits.contiguous(), + lengths_for_kernel, + topk, + persistent_topk_workspace, + topk_tokens, + max_len, + ) + return topk + + fast_topk_v2( + logits.contiguous(), + lengths_for_kernel, + topk, + topk_tokens, + next_n, + ) + return topk + + +def _deepseek_v4_indexer_topk_from_logits_prefill_op( + logits: torch.Tensor, + length_rows: torch.Tensor, + topk_tokens: int, + *, + row_starts: torch.Tensor | None = None, + row_ends: torch.Tensor | None = None, + out: torch.Tensor, +) -> torch.Tensor: + """Use the local TRT-LLM CUDA prefill selector.""" + + if not logits.is_cuda or logits.dtype != torch.float32: + raise RuntimeError("DeepSeek V4 prefill indexer requires CUDA float32 logits") + trtllm_ops = getattr(torch.ops, "trtllm", None) + if trtllm_ops is None or not hasattr(trtllm_ops, "indexer_topk_prefill"): + raise RuntimeError( + "DeepSeek V4 prefill indexer requires the CUDA prefill top-k op" + ) + + num_rows = length_rows.numel() + if num_rows == 0: + return out[:0] + logits = logits.contiguous() + if row_starts is None: + row_starts_for_kernel = torch.zeros( + num_rows, + device=logits.device, + dtype=torch.int32, + ) + else: + row_starts_for_kernel = ( + row_starts.to( + device=logits.device, + dtype=torch.int32, + ) + .reshape(-1) + .contiguous() + ) + if row_ends is None: + row_ends_for_kernel = ( + row_starts_for_kernel + + length_rows.to(device=logits.device, dtype=torch.int32).reshape(-1) + ).contiguous() + else: + row_ends_for_kernel = ( + row_ends.to( + device=logits.device, + dtype=torch.int32, + ) + .reshape(-1) + .contiguous() + ) + + topk = out[:num_rows] + topk.fill_(-1) + trtllm_ops.indexer_topk_prefill( + logits, + row_starts_for_kernel, + row_ends_for_kernel, + topk, + topk_tokens, + ) + return topk + + +@dataclass(frozen=True) +class _DeepseekV4IndexerPrefillChunk: + token_start: int + token_end: int + req_start: int + req_end: int + query_start: int + query_end: int + skip_kv_gather: bool = False + + +def _deepseek_v4_indexer_prefill_max_logits_bytes( + max_logits_bytes: int | None = None, +) -> int: + if max_logits_bytes is not None: + return max(1, int(max_logits_bytes)) + max_logits_mb = global_server_args_dict["deepseek_v4_indexer_prefill_max_logits_mb"] + return max(1, int(max_logits_mb) * 1024 * 1024) + + +def _deepseek_v4_indexer_prefill_workspace_size( + seq_lens_cpu: torch.Tensor, + workspace_size: int | None = None, +) -> int: + if workspace_size is not None: + return max(1, int(workspace_size)) + context_len = global_server_args_dict.get("max_model_len") + if isinstance(context_len, int) and context_len > 0: + return context_len * 40 + max_seq_len = int(seq_lens_cpu.max().item()) if seq_lens_cpu.numel() else 1 + return max(1, max_seq_len) * 40 + + +def _deepseek_v4_indexer_prefill_request_chunks( + *, + seq_lens_cpu: torch.Tensor, + query_lens_cpu: torch.Tensor, + compress_ratio: int, + num_tokens: int, + max_logits_bytes: int | None = None, + workspace_size: int | None = None, + request_offset: int = 0, +) -> list[_DeepseekV4IndexerPrefillChunk]: + """Build request/query-slice sparse-indexer prefill chunks.""" + + if num_tokens == 0: + return [] + + seq_lens = seq_lens_cpu.detach().cpu().to(torch.int64) + query_lens = query_lens_cpu.detach().cpu().to(torch.int64) + if seq_lens.numel() != query_lens.numel(): + return [] + + query_lens_list = [max(0, int(x)) for x in query_lens.tolist()] + if sum(query_lens_list) != num_tokens: + return [] + + compressed_seq_lens = torch.div( + seq_lens, + max(1, int(compress_ratio)), + rounding_mode="floor", + ) + compressed_seq_lens_list = [max(0, int(x)) for x in compressed_seq_lens.tolist()] + workspace_rows = _deepseek_v4_indexer_prefill_workspace_size( + seq_lens, + workspace_size, + ) + max_logits_elems = ( + _deepseek_v4_indexer_prefill_max_logits_bytes(max_logits_bytes) // 4 + ) + max_logits_elems = max(1, max_logits_elems) + + query_offsets = [0] + for query_len in query_lens_list: + query_offsets.append(query_offsets[-1] + query_len) + + chunks: list[_DeepseekV4IndexerPrefillChunk] = [] + n_reqs = len(query_lens_list) + end = 0 + while end < n_reqs: + start = end + chunk_m = 0 + chunk_n = 0 + while end < n_reqs: + q_len = query_lens_list[end] + seq_len = compressed_seq_lens_list[end] + new_m = chunk_m + q_len + new_n = chunk_n + seq_len + if new_n <= workspace_rows and new_m * new_n <= max_logits_elems: + chunk_m = new_m + chunk_n = new_n + end += 1 + else: + break + + if end == start: + chunk_m = query_lens_list[end] + chunk_n = compressed_seq_lens_list[end] + end += 1 + + if chunk_m <= 0: + continue + + req_start = start + request_offset + req_end = end + request_offset + max_q = max(1, max_logits_elems // chunk_n) if chunk_n > 0 else chunk_m + chunk_token_start = query_offsets[start] + for query_start in range(0, chunk_m, max_q): + query_end = min(query_start + max_q, chunk_m) + chunks.append( + _DeepseekV4IndexerPrefillChunk( + token_start=chunk_token_start + query_start, + token_end=chunk_token_start + query_end, + req_start=req_start, + req_end=req_end, + query_start=query_start, + query_end=query_end, + skip_kv_gather=query_start > 0, + ) + ) + return chunks + + +def _deepseek_v4_indexer_decode_max_len( + block_table: torch.Tensor, + cache_block_size: int, + compress_ratio: int, +) -> int: + context_len = global_server_args_dict.get("max_model_len") + if isinstance(context_len, int) and context_len > 0: + return max(1, (context_len + compress_ratio - 1) // compress_ratio) + return max( + 1, + (block_table.shape[1] * cache_block_size + compress_ratio - 1) + // compress_ratio, + ) + + +def _deepseek_v4_indexer_prefill_request_gather_plan( + *, + seq_lens_cpu: torch.Tensor, + query_lens_cpu: torch.Tensor, + block_table: torch.Tensor, + cache_block_size: int, + compress_ratio: int, + req_start: int, + req_end: int, + query_start: int, + query_end: int, + build_slots: bool = True, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]: + device = block_table.device + num_rows = max(0, int(query_end) - int(query_start)) + if num_rows == 0 or req_end <= req_start: + empty_i32 = torch.empty(0, dtype=torch.int32, device=device) + empty_i64 = torch.empty(0, dtype=torch.int64, device=device) + return empty_i64, empty_i32, empty_i32, empty_i32, 0 + + seq_lens_list = ( + seq_lens_cpu.detach().cpu().to(torch.int64)[req_start:req_end].tolist() + ) + query_lens_list = ( + query_lens_cpu.detach().cpu().to(torch.int64)[req_start:req_end].tolist() + ) + if len(seq_lens_list) != len(query_lens_list): + empty_i32 = torch.empty(0, dtype=torch.int32, device=device) + empty_i64 = torch.empty(0, dtype=torch.int64, device=device) + return empty_i64, empty_i32, empty_i32, empty_i32, 0 + + ratio = max(1, int(compress_ratio)) + seq_lens_list = [max(0, int(x)) for x in seq_lens_list] + query_lens_list = [max(0, int(x)) for x in query_lens_list] + compressed_lens_list = [seq_len // ratio for seq_len in seq_lens_list] + total_k = sum(compressed_lens_list) + + query_offsets: list[int] = [0] + for query_len in query_lens_list: + query_offsets.append(query_offsets[-1] + query_len) + + req_local_list: list[int] = [] + row_lens_list: list[int] = [] + req_local = 0 + last_req = max(0, len(query_lens_list) - 1) + for row_offset in range(int(query_start), int(query_end)): + while req_local < last_req and row_offset >= query_offsets[req_local + 1]: + req_local += 1 + local_query_offset = row_offset - query_offsets[req_local] + prefix_len = max(0, seq_lens_list[req_local] - query_lens_list[req_local]) + row_lens_list.append((prefix_len + local_query_offset + 1) // ratio) + req_local_list.append(req_local) + max_len = max(row_lens_list) if row_lens_list else 0 + + compressed_lens = torch.tensor( + compressed_lens_list, + dtype=torch.int64, + device=device, + ) + + cu_seq_lens = torch.empty( + compressed_lens.numel() + 1, + dtype=torch.int32, + device=device, + ) + cu_seq_lens[:1] = 0 + torch.cumsum(compressed_lens.to(torch.int32), dim=0, out=cu_seq_lens[1:]) + + req_local_tensor = torch.tensor(req_local_list, dtype=torch.int64, device=device) + row_lens = torch.tensor(row_lens_list, dtype=torch.int32, device=device) + cu_start = cu_seq_lens[:-1][req_local_tensor] + cu_end = cu_start + row_lens + + if total_k <= 0 or not build_slots: + empty_i64 = torch.empty(0, dtype=torch.int64, device=device) + return empty_i64, cu_start, cu_end, row_lens, max_len + + req_ids = torch.repeat_interleave( + torch.arange(req_start, req_end, device=device, dtype=torch.int64), + compressed_lens, + output_size=total_k, + ) + req_local_for_k = req_ids - int(req_start) + group_bases = cu_seq_lens[:-1][req_local_for_k].to(torch.int64) + local = torch.arange(total_k, device=device, dtype=torch.int64) - group_bases + pages = torch.div(local, cache_block_size, rounding_mode="floor") + page_offsets = local % cache_block_size + page_ids = block_table[req_ids, pages.long()].to(torch.int64) + slots = page_ids * cache_block_size + page_offsets + return slots, cu_start, cu_end, row_lens, max_len + + +def _deepseek_v4_indexer_prefill_chunk_total_rows( + *, + seq_lens_cpu: torch.Tensor, + compress_ratio: int, + req_start: int, + req_end: int, +) -> int: + ratio = max(1, int(compress_ratio)) + seq_lens = seq_lens_cpu.detach().cpu().to(torch.int64)[req_start:req_end].tolist() + return sum(max(0, int(seq_len)) // ratio for seq_len in seq_lens) + + +def _deepseek_v4_indexer_prefill_metadata( + *, + metadata: DeepseekV4ForwardMetadata, + block_table: torch.Tensor, + cache_block_size: int, + compress_ratio: int, + num_prefill_tokens: int, +) -> DeepseekV4IndexerPrefillMetadata: + device = block_table.device + if num_prefill_tokens <= 0: + return DeepseekV4IndexerPrefillMetadata.empty(device) + + seq_lens_cpu = getattr(metadata, "seq_lens_cpu", None) + query_lens_cpu = getattr(metadata, "query_lens_cpu", None) + num_prefill_reqs = int(getattr(metadata, "num_prefill_reqs", 0) or 0) + if seq_lens_cpu is None or query_lens_cpu is None or num_prefill_reqs <= 0: + return DeepseekV4IndexerPrefillMetadata.empty(device) + + seq_lens_cpu = seq_lens_cpu[:num_prefill_reqs] + query_lens_cpu = query_lens_cpu[:num_prefill_reqs] + cache_key = (compress_ratio, cache_block_size, num_prefill_tokens) + cache = metadata.indexer.prefill_plan_cache + cached = cache.get(cache_key) + if cached is not None and cached.slots.device == device: + return cached + + chunks = _deepseek_v4_indexer_prefill_request_chunks( + seq_lens_cpu=seq_lens_cpu, + query_lens_cpu=query_lens_cpu, + compress_ratio=compress_ratio, + num_tokens=num_prefill_tokens, + ) + if not chunks: + out = DeepseekV4IndexerPrefillMetadata.empty(device) + cache[cache_key] = out + return out + + chunk_plans: list[DeepseekV4IndexerPrefillChunkPlan] = [] + slot_parts: list[torch.Tensor] = [] + cu_seq_lens_parts: list[torch.Tensor] = [] + cu_seqlen_k_start_parts: list[torch.Tensor] = [] + cu_seqlen_k_end_parts: list[torch.Tensor] = [] + seq_lens_k_parts: list[torch.Tensor] = [] + slot_offset = 0 + cu_seq_offset = 0 + row_offset = 0 + for chunk in chunks: + slots, cu_seqlen_k_start, cu_seqlen_k_end, seq_lens_k, max_seqlen_k = ( + _deepseek_v4_indexer_prefill_request_gather_plan( + seq_lens_cpu=seq_lens_cpu, + query_lens_cpu=query_lens_cpu, + block_table=block_table, + cache_block_size=cache_block_size, + compress_ratio=compress_ratio, + req_start=chunk.req_start, + req_end=chunk.req_end, + query_start=chunk.query_start, + query_end=chunk.query_end, + build_slots=False, + ) + ) + slot_count = _deepseek_v4_indexer_prefill_chunk_total_rows( + seq_lens_cpu=seq_lens_cpu, + compress_ratio=compress_ratio, + req_start=chunk.req_start, + req_end=chunk.req_end, + ) + compressed_lens = torch.div( + seq_lens_cpu[chunk.req_start : chunk.req_end].to( + dtype=torch.int32, + device=device, + ), + max(1, int(compress_ratio)), + rounding_mode="floor", + ) + cu_seq_lens = torch.empty( + compressed_lens.numel() + 1, + dtype=torch.int32, + device=device, + ) + cu_seq_lens[:1] = 0 + torch.cumsum(compressed_lens, dim=0, out=cu_seq_lens[1:]) + slot_end = slot_offset + slot_count + cu_seq_end = cu_seq_offset + cu_seq_lens.numel() + row_end = row_offset + seq_lens_k.numel() + chunk_plans.append( + DeepseekV4IndexerPrefillChunkPlan( + token_start=chunk.token_start, + token_end=chunk.token_end, + request_start=chunk.req_start, + request_end=chunk.req_end, + slot_start=slot_offset, + slot_end=slot_end, + gather_row_start=row_offset, + gather_row_end=row_end, + max_seq_len_k=max_seqlen_k, + cu_seq_lens_start=cu_seq_offset, + cu_seq_lens_end=cu_seq_end, + skip_kv_gather=chunk.skip_kv_gather, + ) + ) + if slots.numel() > 0: + slot_parts.append(slots) + cu_seq_lens_parts.append(cu_seq_lens) + cu_seqlen_k_start_parts.append(cu_seqlen_k_start) + cu_seqlen_k_end_parts.append(cu_seqlen_k_end) + seq_lens_k_parts.append(seq_lens_k) + slot_offset = slot_end + cu_seq_offset = cu_seq_end + row_offset = row_end + + out = DeepseekV4IndexerPrefillMetadata( + chunks=tuple(chunk_plans), + chunk_specs=torch.tensor( + [ + [ + chunk.token_start, + chunk.token_end, + chunk.request_start, + chunk.request_end, + 1 if chunk.skip_kv_gather else 0, + ] + for chunk in chunk_plans + ], + dtype=torch.int64, + device="cpu", + ), + chunk_offsets=torch.tensor( + [ + [ + chunk.slot_start, + chunk.slot_end, + chunk.gather_row_start, + chunk.gather_row_end, + chunk.max_seq_len_k, + chunk.cu_seq_lens_start, + chunk.cu_seq_lens_end, + ] + for chunk in chunk_plans + ], + dtype=torch.int64, + device="cpu", + ), + slots=( + torch.cat(slot_parts, dim=0) + if slot_parts + else torch.empty(0, dtype=torch.int64, device=device) + ), + cu_seq_lens=( + torch.cat(cu_seq_lens_parts, dim=0) + if cu_seq_lens_parts + else torch.empty(0, dtype=torch.int32, device=device) + ), + cu_seqlen_k_start=( + torch.cat(cu_seqlen_k_start_parts, dim=0) + if cu_seqlen_k_start_parts + else torch.empty(0, dtype=torch.int32, device=device) + ), + cu_seqlen_k_end=( + torch.cat(cu_seqlen_k_end_parts, dim=0) + if cu_seqlen_k_end_parts + else torch.empty(0, dtype=torch.int32, device=device) + ), + seq_lens_k=( + torch.cat(seq_lens_k_parts, dim=0) + if seq_lens_k_parts + else torch.empty(0, dtype=torch.int32, device=device) + ), + ) + cache[cache_key] = out + return out + + +def _deepseek_v4_indexer_decode_plan( + *, + positions: torch.Tensor, + token_to_req_indices: torch.Tensor, + block_table: torch.Tensor, + cache_block_size: int, + compress_ratio: int, + metadata: DeepseekV4ForwardMetadata | None = None, + is_valid_token: torch.Tensor | None = None, + block_table_base_offsets: torch.Tensor | None = None, +) -> DeepseekV4IndexerDecodePlan: + num_tokens = positions.numel() + key = (int(compress_ratio), int(cache_block_size), int(num_tokens)) + indexer_metadata = metadata.indexer if metadata is not None else None + cache = None if indexer_metadata is None else indexer_metadata.decode_plan_cache + refreshed_keys = ( + None + if indexer_metadata is None + else indexer_metadata.decode_plan_refreshed_keys + ) + cached = cache.get(key) if cache is not None else None + # Hot path: the attention metadata builder hook pre-builds the plan tensors + # before per-layer parallel work so indexer layers only read cached buffers. + if cached is not None and refreshed_keys is not None and key in refreshed_keys: + return cached + + if num_tokens == 0: + context_lens = torch.empty((0, 1), dtype=torch.int32, device=positions.device) + block_tables = torch.empty( + (0, 1), + dtype=torch.int32, + device=block_table.device, + ) + plan = DeepseekV4IndexerDecodePlan(context_lens, block_tables, 0) + if cache is not None: + cache[key] = plan + if refreshed_keys is not None: + refreshed_keys.add(key) + return plan + + rows = int(block_table.shape[0]) if block_table.ndim >= 1 else 0 + cols = int(block_table.shape[1]) if block_table.ndim >= 2 else 0 + max_len = _deepseek_v4_indexer_decode_max_len( + block_table, + cache_block_size, + compress_ratio, + ) + max_blocks = max(1, (max_len + cache_block_size - 1) // cache_block_size) + + expected_context_shape = (num_tokens, 1) + expected_block_shape = (num_tokens, max_blocks) + if ( + cached is None + or cached.context_lens.shape != expected_context_shape + or cached.context_lens.device != positions.device + or cached.context_lens.dtype != torch.int32 + or cached.block_table.shape != expected_block_shape + or cached.block_table.device != block_table.device + or cached.block_table.dtype != torch.int32 + ): + context_lens = torch.empty( + expected_context_shape, + dtype=torch.int32, + device=positions.device, + ) + block_tables = torch.empty( + expected_block_shape, + dtype=torch.int32, + device=block_table.device, + ) + plan = DeepseekV4IndexerDecodePlan( + context_lens=context_lens, + block_table=block_tables, + max_context_len=max_len, + ) + if cache is not None: + cache[key] = plan + else: + plan = cached + plan.max_context_len = max_len + + if rows <= 0 or cols <= 0: + plan.context_lens.zero_() + plan.block_table.zero_() + plan.max_context_len = 0 + else: + deepseek_v4_indexer_decode_metadata_compute( + positions=positions, + token_to_req_indices=token_to_req_indices, + block_table=block_table, + cache_block_size=cache_block_size, + compress_ratio=compress_ratio, + max_blocks=max_blocks, + out_context_lens=plan.context_lens, + out_block_tables=plan.block_table, + block_table_base_offsets=block_table_base_offsets, + ) + if is_valid_token is None: + is_valid_token = getattr(metadata, "is_valid_token", None) + if is_valid_token is not None: + valid = is_valid_token[:num_tokens].to( + device=plan.context_lens.device, + dtype=torch.bool, + ) + with torch.inference_mode(): + plan.context_lens.masked_fill_(~valid.view(num_tokens, 1), 0) + plan.block_table.masked_fill_( + ~valid.to(device=plan.block_table.device).view(num_tokens, 1), + 0, + ) + if refreshed_keys is not None: + refreshed_keys.add(key) + return plan + + +def _deepseek_v4_indexer_decode_schedule_metadata( + *, + positions: torch.Tensor, + cache_block_size: int, + compress_ratio: int, + metadata: DeepseekV4ForwardMetadata | None, + context_lens: torch.Tensor | None = None, +) -> torch.Tensor | None: + if positions.numel() == 0: + return None + if deep_gemm is None: + return None + + num_tokens = positions.numel() + if context_lens is None: + compressed_lens = torch.div( + positions.to(torch.int64) + 1, + compress_ratio, + rounding_mode="floor", + ).clamp_min(0) + context_lens = compressed_lens.to(torch.int32).view(num_tokens, 1).contiguous() + schedule_key = (compress_ratio, cache_block_size, num_tokens) + indexer_metadata = metadata.indexer if metadata is not None else None + schedule_cache = ( + None + if indexer_metadata is None + else indexer_metadata.decode_schedule_metadata_cache + ) + schedule_metadata = ( + schedule_cache.get(schedule_key) if schedule_cache is not None else None + ) + + with nvtx_range("indexer_decode_schedule_metadata"): + refreshed = deep_gemm.get_paged_mqa_logits_metadata( + context_lens, + cache_block_size, + deep_gemm.get_num_sms(), + ) + if schedule_metadata is not None: + if ( + schedule_metadata.shape == refreshed.shape + and schedule_metadata.device == refreshed.device + and schedule_metadata.dtype == refreshed.dtype + ): + with torch.inference_mode(): + schedule_metadata.copy_(refreshed) + return schedule_metadata + if schedule_cache is not None: + schedule_cache[schedule_key] = refreshed + return refreshed + schedule_metadata = refreshed + if schedule_cache is not None: + schedule_cache[schedule_key] = schedule_metadata + return schedule_metadata + + +def _deepseek_v4_indexer_topk_prefill_deepgemm( + *, + cache_2d: torch.Tensor, + block_table: torch.Tensor, + cu_seq_lens: torch.Tensor, + cu_start: torch.Tensor, + cu_end: torch.Tensor, + row_lens: torch.Tensor, + max_len: int, + index_q: tuple[torch.Tensor, torch.Tensor], + weights: torch.Tensor, + cache_block_size: int, + topk_tokens: int, + use_prefill_topk_op: bool, + gathered_k: tuple[torch.Tensor, torch.Tensor] | None = None, + gather_workspace: tuple[torch.Tensor, torch.Tensor] | None = None, +) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor] | None]: + q_values, q_scales = index_q + if not _deepseek_v4_deepgemm_fp4_indexer_available(q_values): + raise RuntimeError("DeepSeek V4 sparse indexer requires DeepGEMM FP4 support") + + num_tokens = q_values.shape[0] + if num_tokens == 0: + return ( + torch.empty( + (0, topk_tokens), + device=q_values.device, + dtype=torch.int32, + ), + gathered_k, + ) + if max_len <= 0: + return ( + torch.full( + (num_tokens, topk_tokens), + -1, + device=q_values.device, + dtype=torch.int32, + ), + gathered_k, + ) + + if gathered_k is None: + with nvtx_range("indexer_topk_prefill_gather_paged_mxfp4"): + gathered_k = _deepseek_v4_gather_paged_indexer_mxfp4_cache( + cache_2d, + block_table, + cu_seq_lens, + cache_block_size, + out=gather_workspace, + ) + k_values, k_scales = gathered_k + + with nvtx_range("indexer_topk_prefill_deepgemm_logits"): + logits = deep_gemm.fp8_fp4_mqa_logits( + q=(q_values.contiguous().view(torch.int8), q_scales.contiguous()), + kv=(k_values.contiguous(), k_scales.contiguous()), + weights=weights.contiguous(), + cu_seq_len_k_start=cu_start, + cu_seq_len_k_end=cu_end, + clean_logits=False, + max_seqlen_k=max_len, + logits_dtype=torch.float32, + ) + + with nvtx_range("indexer_topk_prefill_select"): + return ( + _deepseek_v4_indexer_topk_from_logits( + logits, + row_lens, + topk_tokens, + use_prefill_topk_op=use_prefill_topk_op, + ), + gathered_k, + ) + + +def _deepseek_v4_indexer_topk_from_cache_deepgemm_decode( + *, + cache_2d: torch.Tensor, + positions: torch.Tensor, + token_to_req_indices: torch.Tensor, + block_table: torch.Tensor, + cache_block_size: int, + index_q: tuple[torch.Tensor, torch.Tensor], + weights: torch.Tensor, + compress_ratio: int, + topk_tokens: int, + metadata: DeepseekV4ForwardMetadata | None = None, + schedule_metadata: torch.Tensor | None = None, + decode_context_lens: torch.Tensor | None = None, + decode_block_table: torch.Tensor | None = None, + decode_max_context_len: int | None = None, + is_valid_token: torch.Tensor | None = None, + out: torch.Tensor | None = None, + persistent_topk_workspace: torch.Tensor | None = None, +) -> torch.Tensor: + q_values, q_scales = index_q + if not _deepseek_v4_deepgemm_fp4_indexer_available(q_values): + raise RuntimeError("DeepSeek V4 sparse indexer requires DeepGEMM FP4 support") + + num_tokens = positions.numel() + if num_tokens == 0: + if out is not None: + return out[:0] + return torch.empty((0, topk_tokens), device=positions.device, dtype=torch.int32) + if decode_context_lens is not None and decode_block_table is not None: + context_lens = decode_context_lens + block_tables = decode_block_table + max_len = ( + int(decode_max_context_len) + if decode_max_context_len is not None + else int(context_lens.max().item()) + ) + else: + decode_plan = _deepseek_v4_indexer_decode_plan( + positions=positions, + token_to_req_indices=token_to_req_indices, + block_table=block_table, + cache_block_size=cache_block_size, + compress_ratio=compress_ratio, + metadata=metadata, + is_valid_token=is_valid_token, + ) + context_lens = decode_plan.context_lens + block_tables = decode_plan.block_table + max_len = decode_plan.max_context_len + topk = ( + torch.empty( + (num_tokens, topk_tokens), + device=positions.device, + dtype=torch.int32, + ) + if out is None + else out[:num_tokens] + ) + if max_len <= 0: + topk.fill_(-1) + return topk + kv_cache = _deepseek_v4_indexer_mxfp4_cache_view(cache_2d, cache_block_size) + schedule_key = (compress_ratio, cache_block_size, num_tokens) + indexer_metadata = metadata.indexer if metadata is not None else None + schedule_cache = ( + None + if indexer_metadata is None + else indexer_metadata.decode_schedule_metadata_cache + ) + if schedule_metadata is None: + schedule_metadata = ( + schedule_cache.get(schedule_key) if schedule_cache is not None else None + ) + if schedule_metadata is None: + with nvtx_range("indexer_decode_schedule_metadata"): + schedule_metadata = deep_gemm.get_paged_mqa_logits_metadata( + context_lens, + cache_block_size, + deep_gemm.get_num_sms(), + ) + if schedule_cache is not None: + schedule_cache[schedule_key] = schedule_metadata + + with nvtx_range("indexer_decode_deepgemm_logits"): + logits = deep_gemm.fp8_fp4_paged_mqa_logits( + q=( + q_values.contiguous().view(torch.int8).unsqueeze(1), + q_scales.contiguous().unsqueeze(1), + ), + kv_cache=kv_cache, + weights=weights.contiguous(), + context_lens=context_lens, + block_table=block_tables, + schedule_meta=schedule_metadata, + max_context_len=max_len, + clean_logits=False, + logits_dtype=torch.float32, + ) + + with nvtx_range("indexer_decode_topk"): + return _deepseek_v4_indexer_topk_from_logits( + logits, + context_lens, + topk_tokens, + next_n=1, + out=out, + persistent_topk_workspace=persistent_topk_workspace, + ) + + +def _deepseek_v4_sparse_attn_indexer_native( + *, + cache_2d: torch.Tensor, + positions: torch.Tensor, + token_to_req_indices: torch.Tensor, + block_table: torch.Tensor, + seq_lens_cpu: torch.Tensor, + query_lens_cpu: torch.Tensor, + prefill_chunk_specs: torch.Tensor, + prefill_chunk_offsets: torch.Tensor, + prefill_slots: torch.Tensor, + prefill_cu_seq_lens: torch.Tensor, + prefill_cu_seqlen_k_start: torch.Tensor, + prefill_cu_seqlen_k_end: torch.Tensor, + prefill_seq_lens_k: torch.Tensor, + packed_q_values: torch.Tensor, + packed_q_scales: torch.Tensor, + packed_weights: torch.Tensor, + decode_schedule_metadata: torch.Tensor | None, + decode_context_lens: torch.Tensor | None, + decode_block_table: torch.Tensor | None, + decode_max_context_len: int, + topk_indices_buffer: torch.Tensor, + prefill_gather_values_workspace: torch.Tensor, + prefill_gather_scales_workspace: torch.Tensor, + persistent_topk_workspace: torch.Tensor, + cache_block_size: int, + compress_ratio: int, + topk_tokens: int, + num_prefill_tokens: int, + num_decode_tokens: int, +) -> torch.Tensor: + total_tokens = positions.numel() + topk_out = topk_indices_buffer[:total_tokens] + topk_out.fill_(-1) + if total_tokens == 0: + return topk_out + + def fill_prefill() -> None: + if num_prefill_tokens <= 0: + return + + if prefill_chunk_specs.numel() == 0: + raise RuntimeError( + "DeepSeek V4 sparse indexer prefill requires prepared chunk metadata" + ) + + gather_cache_key = None + gathered_k = None + num_chunks = prefill_chunk_specs.shape[0] + for chunk_idx in range(num_chunks): + ( + token_start, + token_end, + req_start, + req_end, + skip_kv_gather_raw, + ) = prefill_chunk_specs[chunk_idx].tolist() + ( + slot_start, + slot_end, + row_start, + row_end, + max_seqlen_k, + cu_seq_start, + cu_seq_end, + ) = prefill_chunk_offsets[chunk_idx].tolist() + skip_kv_gather = bool(int(skip_kv_gather_raw)) + gather_rows = max(0, slot_end - slot_start) + gather_workspace = None + if ( + prefill_gather_values_workspace.numel() > 0 + and prefill_gather_scales_workspace.numel() > 0 + and gather_rows <= prefill_gather_values_workspace.shape[0] + and gather_rows <= prefill_gather_scales_workspace.shape[0] + ): + gather_workspace = ( + prefill_gather_values_workspace[:gather_rows], + prefill_gather_scales_workspace[:gather_rows], + ) + with nvtx_range("indexer_topk_deepgemm_prefill"): + key = (req_start, req_end) + reuse_k = None + if skip_kv_gather and gather_cache_key == key: + reuse_k = gathered_k + if prefill_cu_seq_lens.numel() == 0 or cu_seq_end <= cu_seq_start: + raise RuntimeError( + "DeepSeek V4 sparse indexer prefill metadata is incomplete" + ) + topk, next_gathered_k = _deepseek_v4_indexer_topk_prefill_deepgemm( + cache_2d=cache_2d, + block_table=block_table[req_start:req_end], + cu_seq_lens=prefill_cu_seq_lens[cu_seq_start:cu_seq_end], + cu_start=prefill_cu_seqlen_k_start[row_start:row_end], + cu_end=prefill_cu_seqlen_k_end[row_start:row_end], + row_lens=prefill_seq_lens_k[row_start:row_end], + max_len=max_seqlen_k, + cache_block_size=cache_block_size, + index_q=( + packed_q_values[token_start:token_end], + packed_q_scales[token_start:token_end], + ), + weights=packed_weights[token_start:token_end], + topk_tokens=topk_tokens, + use_prefill_topk_op=True, + gathered_k=reuse_k, + gather_workspace=gather_workspace, + ) + if next_gathered_k is not None: + gather_cache_key = key + gathered_k = next_gathered_k + topk_out[token_start:token_end].copy_(topk) + + def fill_decode() -> None: + if num_decode_tokens <= 0: + return + + decode_start = num_prefill_tokens + decode_end = decode_start + num_decode_tokens + decode_positions = positions[decode_start:decode_end] + decode_token_to_req = token_to_req_indices[decode_start:decode_end] + decode_out = topk_out[decode_start:decode_end] + with nvtx_range("indexer_topk_deepgemm_decode"): + _deepseek_v4_indexer_topk_from_cache_deepgemm_decode( + cache_2d=cache_2d, + positions=decode_positions, + token_to_req_indices=decode_token_to_req, + block_table=block_table, + cache_block_size=cache_block_size, + index_q=( + packed_q_values[decode_start:decode_end], + packed_q_scales[decode_start:decode_end], + ), + weights=packed_weights[decode_start:decode_end], + compress_ratio=compress_ratio, + topk_tokens=topk_tokens, + schedule_metadata=decode_schedule_metadata, + decode_context_lens=decode_context_lens, + decode_block_table=decode_block_table, + decode_max_context_len=decode_max_context_len, + out=decode_out, + persistent_topk_workspace=persistent_topk_workspace, + ) + + fill_prefill() + fill_decode() + return topk_out + + +def _deepseek_v4_sparse_attn_indexer_op( + cache_2d: torch.Tensor, + positions: torch.Tensor, + token_to_req_indices: torch.Tensor, + block_table: torch.Tensor, + seq_lens_cpu: torch.Tensor, + query_lens_cpu: torch.Tensor, + prefill_chunk_specs: torch.Tensor, + prefill_chunk_offsets: torch.Tensor, + prefill_slots: torch.Tensor, + prefill_cu_seq_lens: torch.Tensor, + prefill_cu_seqlen_k_start: torch.Tensor, + prefill_cu_seqlen_k_end: torch.Tensor, + prefill_seq_lens_k: torch.Tensor, + packed_q_values: torch.Tensor, + packed_q_scales: torch.Tensor, + packed_weights: torch.Tensor, + decode_schedule_metadata: torch.Tensor, + decode_context_lens: torch.Tensor, + decode_block_table: torch.Tensor, + decode_max_context_len: int, + topk_indices_buffer: torch.Tensor, + prefill_gather_values_workspace: torch.Tensor, + prefill_gather_scales_workspace: torch.Tensor, + persistent_topk_workspace: torch.Tensor, + cache_block_size: int, + compress_ratio: int, + topk_tokens: int, + num_prefill_tokens: int, + num_decode_tokens: int, +) -> torch.Tensor: + schedule_metadata = ( + decode_schedule_metadata if decode_schedule_metadata.numel() > 0 else None + ) + context_lens = decode_context_lens if decode_context_lens.numel() > 0 else None + decode_blocks = decode_block_table if decode_block_table.numel() > 0 else None + return _deepseek_v4_sparse_attn_indexer_native( + cache_2d=cache_2d, + positions=positions, + token_to_req_indices=token_to_req_indices, + block_table=block_table, + seq_lens_cpu=seq_lens_cpu, + query_lens_cpu=query_lens_cpu, + prefill_chunk_specs=prefill_chunk_specs, + prefill_chunk_offsets=prefill_chunk_offsets, + prefill_slots=prefill_slots, + prefill_cu_seq_lens=prefill_cu_seq_lens, + prefill_cu_seqlen_k_start=prefill_cu_seqlen_k_start, + prefill_cu_seqlen_k_end=prefill_cu_seqlen_k_end, + prefill_seq_lens_k=prefill_seq_lens_k, + packed_q_values=packed_q_values, + packed_q_scales=packed_q_scales, + packed_weights=packed_weights, + decode_schedule_metadata=schedule_metadata, + decode_context_lens=context_lens, + decode_block_table=decode_blocks, + decode_max_context_len=decode_max_context_len, + topk_indices_buffer=topk_indices_buffer, + prefill_gather_values_workspace=prefill_gather_values_workspace, + prefill_gather_scales_workspace=prefill_gather_scales_workspace, + persistent_topk_workspace=persistent_topk_workspace, + cache_block_size=cache_block_size, + compress_ratio=compress_ratio, + topk_tokens=topk_tokens, + num_prefill_tokens=num_prefill_tokens, + num_decode_tokens=num_decode_tokens, + ) + + +def _deepseek_v4_sparse_attn_indexer_fake( + cache_2d: torch.Tensor, + positions: torch.Tensor, + token_to_req_indices: torch.Tensor, + block_table: torch.Tensor, + seq_lens_cpu: torch.Tensor, + query_lens_cpu: torch.Tensor, + prefill_chunk_specs: torch.Tensor, + prefill_chunk_offsets: torch.Tensor, + prefill_slots: torch.Tensor, + prefill_cu_seq_lens: torch.Tensor, + prefill_cu_seqlen_k_start: torch.Tensor, + prefill_cu_seqlen_k_end: torch.Tensor, + prefill_seq_lens_k: torch.Tensor, + packed_q_values: torch.Tensor, + packed_q_scales: torch.Tensor, + packed_weights: torch.Tensor, + decode_schedule_metadata: torch.Tensor, + decode_context_lens: torch.Tensor, + decode_block_table: torch.Tensor, + decode_max_context_len: int, + topk_indices_buffer: torch.Tensor, + prefill_gather_values_workspace: torch.Tensor, + prefill_gather_scales_workspace: torch.Tensor, + persistent_topk_workspace: torch.Tensor, + cache_block_size: int, + compress_ratio: int, + topk_tokens: int, + num_prefill_tokens: int, + num_decode_tokens: int, +) -> torch.Tensor: + del ( + cache_2d, + positions, + token_to_req_indices, + block_table, + seq_lens_cpu, + query_lens_cpu, + prefill_chunk_specs, + prefill_chunk_offsets, + prefill_slots, + prefill_cu_seq_lens, + prefill_cu_seqlen_k_start, + prefill_cu_seqlen_k_end, + prefill_seq_lens_k, + packed_q_values, + packed_q_scales, + packed_weights, + decode_schedule_metadata, + decode_context_lens, + decode_block_table, + decode_max_context_len, + cache_block_size, + prefill_gather_values_workspace, + prefill_gather_scales_workspace, + persistent_topk_workspace, + compress_ratio, + topk_tokens, + num_prefill_tokens, + num_decode_tokens, + ) + return topk_indices_buffer + + +direct_register_custom_op( + op_name="deepseek_v4_sparse_attn_indexer", + op_func=_deepseek_v4_sparse_attn_indexer_op, + mutates_args=[ + "topk_indices_buffer", + "prefill_gather_values_workspace", + "prefill_gather_scales_workspace", + "persistent_topk_workspace", + ], + fake_impl=_deepseek_v4_sparse_attn_indexer_fake, +) + + +def _deepseek_v4_sparse_attn_indexer( + *, + indexer_metadata: DeepseekV4SparseIndexerMetadata, + indexer_cache: torch.Tensor, + indexer_block_table: torch.Tensor, + indexer_block_size: int, + compress_ratio: int, + packed_q_values: torch.Tensor, + packed_q_scales: torch.Tensor, + packed_weights: torch.Tensor, + topk_indices_buffer: torch.Tensor, + prefill_gather_values_workspace: torch.Tensor, + prefill_gather_scales_workspace: torch.Tensor, + persistent_topk_workspace: torch.Tensor, + topk_tokens: int, +) -> torch.Tensor: + batch_metadata = indexer_metadata.batch_metadata + prefill_metadata = indexer_metadata.prefill_metadata + if batch_metadata is None or prefill_metadata is None: + raise RuntimeError("DeepSeek V4 sparse indexer metadata is not prepared") + + positions = batch_metadata.positions + decode_plan = indexer_metadata.decode_plan + decode_schedule_metadata = indexer_metadata.decode_schedule_metadata + decode_context_lens = None if decode_plan is None else decode_plan.context_lens + decode_block_table = None if decode_plan is None else decode_plan.block_table + decode_max_context_len = 0 if decode_plan is None else decode_plan.max_context_len + if decode_schedule_metadata is None: + decode_schedule_metadata = torch.empty( + 0, + dtype=torch.int32, + device=positions.device, + ) + if decode_context_lens is None: + decode_context_lens = torch.empty( + (0, 1), + dtype=torch.int32, + device=positions.device, + ) + if decode_block_table is None: + decode_block_table = torch.empty( + (0, 1), + dtype=indexer_block_table.dtype, + device=indexer_block_table.device, + ) + if not positions.is_cuda: + raise RuntimeError("DeepSeek V4 sparse indexer requires CUDA tensors") + return torch.ops.tokenspeed.deepseek_v4_sparse_attn_indexer( + indexer_cache, + positions, + batch_metadata.token_to_req_indices, + indexer_block_table, + batch_metadata.seq_lens_cpu, + batch_metadata.query_lens_cpu, + prefill_metadata.chunk_specs, + prefill_metadata.chunk_offsets, + prefill_metadata.slots, + prefill_metadata.cu_seq_lens, + prefill_metadata.cu_seqlen_k_start, + prefill_metadata.cu_seqlen_k_end, + prefill_metadata.seq_lens_k, + packed_q_values, + packed_q_scales, + packed_weights, + decode_schedule_metadata, + decode_context_lens, + decode_block_table, + decode_max_context_len, + topk_indices_buffer, + prefill_gather_values_workspace, + prefill_gather_scales_workspace, + persistent_topk_workspace, + indexer_block_size, + compress_ratio, + topk_tokens, + batch_metadata.num_prefill_tokens, + batch_metadata.num_decode_tokens, + ) + + +def _deepseek_v4_mega_moe_max_num_tokens() -> int: + override = int( + global_server_args_dict.get("deepseek_v4_mega_moe_max_num_tokens", 0) or 0 + ) + if override > 0: + return override + + candidates = [ + global_server_args_dict.get("chunked_prefill_size", 0), + global_server_args_dict.get("prefill_graph_max_tokens", 0), + global_server_args_dict.get("max_cudagraph_capture_size", 0), + global_server_args_dict.get("max_num_seqs", 0), + ] + return max([int(value or 0) for value in candidates] + [1]) + + +class _DeepseekV4TopKBuffer: + def __init__(self, topk_tokens: int) -> None: + self.topk_tokens = topk_tokens + self.buffer: torch.Tensor | None = None + + def get(self, num_tokens: int, device: torch.device) -> torch.Tensor: + rows = max(num_tokens, _deepseek_v4_mega_moe_max_num_tokens()) + needs_alloc = ( + self.buffer is None + or self.buffer.device != device + or self.buffer.shape[0] < rows + or self.buffer.shape[1] != self.topk_tokens + ) + if needs_alloc: + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "DeepSeek V4 top-k buffer must be allocated before CUDA graph " + "capture" + ) + self.buffer = torch.empty( + (rows, self.topk_tokens), + dtype=torch.int32, + device=device, + ) + return self.buffer[:num_tokens] + + +def _deepseek_v4_padded_heads(num_local_heads: int) -> int: + if num_local_heads <= 64: + return 64 + if num_local_heads <= 128: + return 128 + raise ValueError( + f"DeepSeek V4 attention supports at most 128 local heads, got {num_local_heads}" + ) + + +def _deepseek_v4_sanitize_swa_slot_mapping( + slot_mapping: torch.Tensor, + capacity: int, + is_valid_token: torch.Tensor | None = None, +) -> torch.Tensor: + # Fail closed: with no writable SWA capacity every slot is masked, so an + # unchecked mapping can never reach the fused cache-insert kernels. + if capacity <= 0: + return torch.full_like(slot_mapping, -1) + slot_mapping = _mask_invalid_graph_tokens(slot_mapping, is_valid_token) + valid = (slot_mapping >= 0) & (slot_mapping < capacity) + return torch.where( + valid, + slot_mapping, + torch.full_like(slot_mapping, -1), + ) + + +def _deepseek_v4_swa_slot_mapping( + ctx: ForwardContext, + positions: torch.Tensor, + out_cache_loc: torch.Tensor, +) -> torch.Tensor: + """Build the SWA write-slot mapping, already sanitized for cache inserts. + + The returned mapping has invalid CUDA-graph tokens and out-of-capacity + slots masked to -1, so per-layer SWA inserts can consume it directly. + Sanitizing here keeps the mask/clamp elementwise chain at once per step; + doing it per layer previously baked ~7 tiny kernels x 61 layers into the + captured decode graph. + """ + if positions.numel() == 0: + return out_cache_loc + metadata = _deepseek_v4_forward_metadata(ctx) + if metadata is None: + raise RuntimeError("DeepSeek V4 attention requires forward metadata") + cache_metadata = metadata.cache + token_to_req_indices = metadata.token_to_req_indices[: positions.numel()] + if cache_metadata.swa_block_table is None: + slot_mapping = out_cache_loc + elif token_to_req_indices.numel() != positions.numel() and ( + token_to_req_indices.numel() <= 0 + or positions.numel() % token_to_req_indices.numel() != 0 + ): + slot_mapping = out_cache_loc + else: + slot_mapping = _group_slot_mapping_from_raw( + positions, + token_to_req_indices, + cache_metadata.swa_block_table, + ctx.token_to_kv_pool.swa_block_size, + base_offsets=cache_metadata.swa_base_logical_page, + ) + is_valid_token = getattr(metadata, "is_valid_token", None) + if is_valid_token is not None: + is_valid_token = is_valid_token[: positions.numel()] + # Attribute access is deliberately unguarded: a pool without + # swa_capacity_slots must fail fast here rather than skip the bounds + # check that protects the fused cache-insert kernels. + return _deepseek_v4_sanitize_swa_slot_mapping( + slot_mapping, + ctx.token_to_kv_pool.swa_capacity_slots, + is_valid_token, + ) + + +def _attention_use_fp4_indexer_cache(config: PretrainedConfig) -> bool: + override = global_server_args_dict.get("attention_use_fp4_indexer_cache", None) + if override is not None: + return bool(override) + attention_config = getattr(config, "attention_config", None) + if isinstance(attention_config, dict): + return bool(attention_config.get("use_fp4_indexer_cache", False)) + return bool(getattr(attention_config, "use_fp4_indexer_cache", False)) + + +def deepseek_v4_rope_config( + config: PretrainedConfig, compress_ratio: int +) -> tuple[float, dict | None]: + """Return the per-layer DeepSeek V4 RoPE base and scaling config. + + DeepSeek V4 uses ordinary RoPE for SWA-only layers. Compressed layers + use the checkpoint's separate `compress_rope_theta` together with YaRN. + """ + + if compress_ratio <= 1: + return float(getattr(config, "rope_theta", 10000.0)), None + + rope_scaling = getattr(config, "rope_scaling", None) + if rope_scaling is not None: + rope_scaling = dict(rope_scaling) + rope_scaling["rope_type"] = "deepseek_yarn" + rope_scaling["mscale"] = 0 + rope_scaling["mscale_all_dim"] = 0 + return ( + float( + getattr( + config, + "compress_rope_theta", + getattr(config, "rope_theta", 10000.0), + ) + ), + rope_scaling, + ) + + +class DeepseekV4MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + mapping: Mapping, + quant_config: QuantizationConfig | None, + prefix: str, + swiglu_limit: float | None = None, + reduce_results: bool = False, + ) -> None: + super().__init__() + if hidden_act != "silu": + raise ValueError(f"Unsupported activation: {hidden_act}") + tp = mapping.dense + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + tp_rank=tp.tp_rank, + tp_size=tp.tp_size, + tp_group=tp.tp_group, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + reduce_results=reduce_results, + tp_rank=tp.tp_rank, + tp_size=tp.tp_size, + tp_group=tp.tp_group, + quant_config=quant_config, + prefix=add_prefix("down_proj", prefix), + ) + self.swiglu_limit = swiglu_limit + self.reduce_results = reduce_results + self.tp_rank = tp.tp_rank + self.tp_size = tp.tp_size + self.tp_group = tp.tp_group + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if x.shape[0] == 0: + return x.new_empty((0, self.down_proj.output_size)) + gate_up, _ = self.gate_up_proj(x) + + _use_fused_swiglu = ( + gate_up.ndim == 2 + and gate_up.shape[-1] % 256 == 0 + and getattr(self.down_proj, "_use_deep_gemm_fp8", False) + ) + if _use_fused_swiglu: + from tokenspeed_kernel.ops.activation.triton import ( + fused_swiglu_fp8_ue8m0, + ) + + x_fp8, scale = fused_swiglu_fp8_ue8m0( + gate_up, swiglu_limit=self.swiglu_limit or 0.0 + ) + out, _ = self.down_proj(x_fp8, scale=scale) + else: + gate, up = gate_up.float().chunk(2, dim=-1) + if self.swiglu_limit is not None and self.swiglu_limit > 0: + gate = torch.clamp(gate, max=self.swiglu_limit) + up = torch.clamp(up, min=-self.swiglu_limit, max=self.swiglu_limit) + x = (F.silu(gate) * up).to(x.dtype) + out, _ = self.down_proj(x) + return out + + +class DeepseekV4MoEGate(nn.Module): + def __init__( + self, + config: PretrainedConfig, + layer_index: int, + hash_indices_dtype: torch.dtype = torch.int32, + ) -> None: + super().__init__() + self.weight = nn.Parameter( + torch.empty(config.n_routed_experts, config.hidden_size) + ) + self.is_hash_moe = layer_index < config.num_hash_layers + if self.is_hash_moe: + self.tid2eid = nn.Parameter( + torch.empty( + config.vocab_size, + config.num_experts_per_tok, + dtype=hash_indices_dtype, + ), + requires_grad=False, + ) + self.e_score_correction_bias = None + elif getattr(config, "topk_method", None) == "noaux_tc": + self.register_parameter("tid2eid", None) + self.e_score_correction_bias = nn.Parameter( + torch.empty(config.n_routed_experts, dtype=torch.float32), + requires_grad=False, + ) + else: + self.register_parameter("tid2eid", None) + self.e_score_correction_bias = None + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return _deepseek_v4_router_gemm(hidden_states, self.weight) + + +class DeepseekV4MegaMoEExperts(nn.Module): + _symm_buffer_cache: dict[tuple[int, int, int, int, int, int, int], object] = {} + + def __init__( + self, + *, + num_experts: int, + num_local_experts: int, + top_k: int, + hidden_size: int, + intermediate_size: int, + mapping: Mapping, + prefix: str, + swiglu_limit: float | None, + ) -> None: + super().__init__() + self.prefix = prefix + self.num_experts = num_experts + self.num_local_experts = num_local_experts + self.top_k = top_k + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.mapping = mapping + self.max_num_tokens = _deepseek_v4_mega_moe_max_num_tokens() + # DeepGEMM bakes the activation clamp into the kernel as a compile-time + # template arg, so the warmup below must use the same value serving does + # (the caller passes it to ``forward``) for the pre-compiled tiles to + # match the served kernels. + self.swiglu_limit = swiglu_limit + + weight_attrs = {"weight_loader": self.weight_loader} + self.w13_weight = nn.Parameter( + torch.zeros( + num_local_experts, + 2 * intermediate_size, + hidden_size // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ) + set_weight_attrs(self.w13_weight, weight_attrs) + + self.w13_weight_scale = nn.Parameter( + torch.zeros( + num_local_experts, + 2 * intermediate_size, + hidden_size // DEEPSEEK_V4_MXFP4_BLOCK_SIZE, + dtype=torch.uint8, + ), + requires_grad=False, + ) + set_weight_attrs(self.w13_weight_scale, weight_attrs) + + self.w2_weight = nn.Parameter( + torch.zeros( + num_local_experts, + hidden_size, + intermediate_size // 2, + dtype=torch.uint8, + ), + requires_grad=False, + ) + set_weight_attrs(self.w2_weight, weight_attrs) + + self.w2_weight_scale = nn.Parameter( + torch.zeros( + num_local_experts, + hidden_size, + intermediate_size // DEEPSEEK_V4_MXFP4_BLOCK_SIZE, + dtype=torch.uint8, + ), + requires_grad=False, + ) + set_weight_attrs(self.w2_weight_scale, weight_attrs) + + self._transformed_l1_weights: tuple[torch.Tensor, torch.Tensor] | None = None + self._transformed_l2_weights: tuple[torch.Tensor, torch.Tensor] | None = None + + def weight_loader( + self, + param: nn.Parameter, + loaded_weight: torch.Tensor, + shard_id: str, + local_expert_id: int, + ) -> None: + expert_data = param.data[local_expert_id] + if shard_id in ("w1", "w3"): + if param is not self.w13_weight and param is not self.w13_weight_scale: + raise ValueError(f"Unexpected MegaMoE w13 shard target: {shard_id}") + shard_offset = 0 if shard_id == "w1" else self.intermediate_size + expert_data = expert_data.narrow(0, shard_offset, self.intermediate_size) + elif shard_id == "w2": + if param is not self.w2_weight and param is not self.w2_weight_scale: + raise ValueError(f"Unexpected MegaMoE w2 shard target: {shard_id}") + else: + raise ValueError(f"Unsupported DeepSeek V4 MegaMoE shard id: {shard_id}") + + if expert_data.dtype == torch.uint8 and loaded_weight.dtype == getattr( + torch, "float8_e8m0fnu", None + ): + loaded_weight = loaded_weight.view(torch.uint8) + if expert_data.shape != loaded_weight.shape: + raise ValueError( + f"DeepSeek V4 MegaMoE expert weight shape mismatch for " + f"{self.prefix}: parameter shard {tuple(expert_data.shape)} " + f"vs checkpoint {tuple(loaded_weight.shape)}" + ) + expert_data.copy_(loaded_weight) + + @staticmethod + def _ue8m0_to_float(sf: torch.Tensor) -> torch.Tensor: + if sf.dtype == torch.uint8: + return (sf.to(torch.int32) << 23).view(torch.float32) + return sf.float() + + def _check_runtime_supported(self) -> None: + if not torch.cuda.is_available(): + raise NotImplementedError("DeepSeek V4 MegaMoE requires CUDA.") + device = self.w13_weight.device + if device.type != "cuda": + raise NotImplementedError( + "DeepSeek V4 MegaMoE expert weights must be loaded on CUDA." + ) + if torch.cuda.get_device_capability(device)[0] != 10: + raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.") + if self.hidden_size % 128 != 0 or self.intermediate_size % 128 != 0: + raise ValueError( + "DeepGEMM MegaMoE requires hidden and intermediate sizes " + "to be multiples of 128." + ) + + def finalize_weights(self) -> None: + if self._transformed_l1_weights is not None: + return + + self._check_runtime_supported() + if deep_gemm is None: + raise RuntimeError( + "DeepSeek V4 MegaMoE backend requires a DeepGEMM package with " + "fp8_fp4_mega_moe support; pass --moe-backend mega_moe only " + "when DeepGEMM is installed." + ) + + w13_scale = deep_gemm.transform_sf_into_required_layout( + self._ue8m0_to_float(self.w13_weight_scale.data).contiguous(), + 2 * self.intermediate_size, + self.hidden_size, + (1, DEEPSEEK_V4_MXFP4_BLOCK_SIZE), + self.num_local_experts, + ) + w2_scale = deep_gemm.transform_sf_into_required_layout( + self._ue8m0_to_float(self.w2_weight_scale.data).contiguous(), + self.hidden_size, + self.intermediate_size, + (1, DEEPSEEK_V4_MXFP4_BLOCK_SIZE), + self.num_local_experts, + ) + l1, l2 = deep_gemm.transform_weights_for_mega_moe( + (self.w13_weight.data.view(torch.int8).contiguous(), w13_scale), + (self.w2_weight.data.view(torch.int8).contiguous(), w2_scale), + ) + # L2 data tensor may alias the input .contiguous() buffer — clone + # to break the reference so the original weight storage can be freed. + self._transformed_l1_weights = l1 + self._transformed_l2_weights = (l2[0].clone(), l2[1]) + + del self.w13_weight + del self.w13_weight_scale + del self.w2_weight + del self.w2_weight_scale + + def get_symm_buffer(self): + if deep_gemm is None: + raise RuntimeError("DeepGEMM MegaMoE symbols are unavailable.") + group = pg_manager.get_process_group("nccl", self.mapping.moe.tp_ep_group) + device = torch.cuda.current_device() + key = ( + id(group), + device, + self.num_experts, + self.max_num_tokens, + self.top_k, + self.hidden_size, + self.intermediate_size, + ) + symm_buffer = self._symm_buffer_cache.get(key) + if symm_buffer is None: + symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe( + group, + self.num_experts, + self.max_num_tokens, + self.top_k, + self.hidden_size, + self.intermediate_size, + ) + self._symm_buffer_cache[key] = symm_buffer + return symm_buffer + + def forward( + self, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + *, + activation_clamp: float | None, + fast_math: bool = True, + ) -> torch.Tensor: + if hidden_states.shape[0] > self.max_num_tokens: + raise ValueError( + f"DeepSeek V4 MegaMoE got {hidden_states.shape[0]} tokens, " + f"but the symmetric buffer was sized for {self.max_num_tokens}." + ) + + y = torch.empty_like(hidden_states, dtype=torch.bfloat16) + symm_buffer = self.get_symm_buffer() + num_tokens = hidden_states.shape[0] + _stage_deepseek_v4_mega_moe_inputs( + hidden_states, + topk_weights, + topk_ids, + symm_buffer.x[:num_tokens], + symm_buffer.x_sf[:num_tokens], + symm_buffer.topk_idx[:num_tokens], + symm_buffer.topk_weights[:num_tokens], + ) + + if self._transformed_l1_weights is None or self._transformed_l2_weights is None: + raise RuntimeError( + "DeepseekV4MegaMoEExperts.finalize_weights() must run via " + "post_load_weights() before forward()" + ) + if deep_gemm is None: + raise RuntimeError("deep_gemm is required for fp8_fp4_mega_moe.") + deep_gemm.fp8_fp4_mega_moe( + y, + self._transformed_l1_weights, + self._transformed_l2_weights, + symm_buffer, + activation_clamp=activation_clamp, + fast_math=fast_math, + ) + return y + + +class DeepseekV4MoE(nn.Module): + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None, + layer_index: int, + prefix: str, + aux_stream: torch.cuda.Stream | None = None, + ) -> None: + super().__init__() + self.config = config + self.mapping = mapping + self.layer_index = layer_index + self.n_shared_experts = config.n_shared_experts + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) + self.scoring_func = getattr(config, "scoring_func", "sqrtsoftplus") + if self.scoring_func != "sqrtsoftplus": + raise ValueError( + f"Unsupported DeepSeek V4 MoE scoring: {self.scoring_func}" + ) + self.stream_fork = StreamFork(aux_stream) + + self.use_mega_moe = get_moe_backend().is_mega_moe() + if self.use_mega_moe: + if deep_gemm is None: + raise RuntimeError( + "DeepSeek V4 MegaMoE backend requires an external DeepGEMM " + "package with fp8_fp4_mega_moe support." + ) + if mapping.moe.ep_size <= 1: + raise ValueError("DeepSeek V4 MegaMoE requires expert parallelism.") + if mapping.moe.tp_size != 1: + raise ValueError("DeepSeek V4 MegaMoE does not support mixed TP/EP.") + if global_server_args_dict.get("ep_num_redundant_experts", 0): + raise ValueError( + "DeepSeek V4 MegaMoE does not support redundant EP experts." + ) + if config.n_routed_experts % mapping.moe.ep_size != 0: + raise ValueError( + "DeepSeek V4 MegaMoE requires n_routed_experts divisible by " + "EP size." + ) + self.hash_indices_dtype = torch.int64 if self.use_mega_moe else torch.int32 + self.gate = DeepseekV4MoEGate( + config, + layer_index, + hash_indices_dtype=self.hash_indices_dtype, + ) + + if config.n_shared_experts is not None: + self.shared_experts = DeepseekV4MLP( + config.hidden_size, + config.moe_intermediate_size * config.n_shared_experts, + config.hidden_act, + mapping, + quant_config, + add_prefix("shared_experts", prefix), + swiglu_limit=getattr(config, "swiglu_limit", None), + reduce_results=False, + ) + else: + self.shared_experts = None + + if self.use_mega_moe: + self.experts = DeepseekV4MegaMoEExperts( + num_experts=config.n_routed_experts, + num_local_experts=config.n_routed_experts // mapping.moe.ep_size, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + mapping=mapping, + prefix=add_prefix("experts", prefix), + swiglu_limit=getattr(config, "swiglu_limit", None), + ) + self.topk = None + else: + routed_quant_config = Mxfp4Config( + ignored_layers=quant_config.ignored_layers, + is_checkpoint_mxfp4_serialized=True, + ) + self.experts = MoELayer( + top_k=config.num_experts_per_tok, + num_experts=config.n_routed_experts + + global_server_args_dict["ep_num_redundant_experts"], + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + quant_config=routed_quant_config, + layer_index=layer_index, + prefix=prefix, + tp_rank=mapping.moe.tp_rank, + tp_size=mapping.moe.tp_size, + ep_rank=mapping.moe.ep_rank, + ep_size=mapping.moe.ep_size, + activation="swiglu", + swiglu_limit=getattr(config, "swiglu_limit", None), + with_bias=True, + routing_config={ + "routed_scaling_factor": self.routed_scaling_factor, + "normalize_topk_weights": config.norm_topk_prob, + "correction_bias": self.gate.e_score_correction_bias, + "routing_method_type": RoutingMethodType.Renormalize, + }, + ) + self.topk = TopK( + top_k=config.num_experts_per_tok, + renormalize=config.norm_topk_prob, + correction_bias=self.gate.e_score_correction_bias, + routed_scaling_factor=self.routed_scaling_factor, + output_format=self.experts.topk_output_format, + ) + + def _select_experts( + self, + hidden_states: torch.Tensor, + input_ids: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + router_logits = self.gate(hidden_states) + fmt = getattr(self.experts, "topk_output_format", None) + need_scores = fmt is not None and not fmt.is_bypassed() + return deepseek_v4_select_experts( + router_logits, + self.config.num_experts_per_tok, + self.config.norm_topk_prob, + correction_bias=self.gate.e_score_correction_bias, + hash_indices_table=self.gate.tid2eid, + input_ids=input_ids, + need_scores=need_scores, + ) + + def _make_topk_output( + self, + hidden_states: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + router_scores: torch.Tensor, + ): + if self.experts.topk_output_format.is_bypassed(): + router_logits = pack_topk_as_router_logits( + topk_weights, topk_ids, self.config.n_routed_experts + ) + return BypassedTopKOutput( + hidden_states, router_logits, self.topk.topk_config + ) + return StandardTopKOutput(topk_weights, topk_ids, router_scores) + + def _forward_shared_experts( + self, + hidden_states: torch.Tensor, + ctx: ForwardContext | None = None, + comm_manager: CommManager | None = None, + ) -> torch.Tensor | None: + if self.shared_experts is None: + return None + if comm_manager is not None: + hidden_states = comm_manager.pre_dense_comm(hidden_states, ctx) + if hidden_states.shape[0] == 0: + return None + with nvtx_range("moe_shared_experts"): + shared = self.shared_experts(hidden_states) + if comm_manager is not None: + shared, _ = comm_manager.post_dense_comm(shared, None, ctx) + return shared + + def forward_mega_moe( + self, + hidden_states: torch.Tensor, + input_ids: torch.Tensor, + ctx: ForwardContext, + comm_manager: CommManager, + ) -> torch.Tensor: + if hidden_states.shape[0] == 0: + topk_weights = hidden_states.new_empty( + (0, self.config.num_experts_per_tok), dtype=torch.float32 + ) + topk_ids = torch.empty( + (0, self.config.num_experts_per_tok), + device=hidden_states.device, + dtype=torch.int64, + ) + else: + with nvtx_range("moe_select_experts"): + topk_weights, topk_ids, _ = self._select_experts( + hidden_states, input_ids + ) + + shared = None + with self.stream_fork.scope(enable=get_is_capture_mode()) as fork: + with nvtx_range("moe_mega_experts"): + if topk_ids.dtype != torch.int64: + topk_ids = topk_ids.to(torch.int64) + if self.routed_scaling_factor != 1.0: + topk_weights = topk_weights * self.routed_scaling_factor + routed = self.experts( + hidden_states, + topk_weights, + topk_ids, + activation_clamp=getattr(self.config, "swiglu_limit", None), + ) + with fork.branch(): + shared = self._forward_shared_experts( + hidden_states, + ctx=ctx, + comm_manager=comm_manager, + ) + return routed + shared if shared is not None else routed + + def forward_normal( + self, + hidden_states: torch.Tensor, + input_ids: torch.Tensor, + num_global_tokens: int, + max_num_tokens_per_gpu: int, + ) -> torch.Tensor: + if hidden_states.shape[0] == 0: + return hidden_states + with nvtx_range("moe_select_experts"): + topk_weights, topk_ids, router_scores = self._select_experts( + hidden_states, input_ids + ) + with nvtx_range("moe_make_topk_output"): + topk_output = self._make_topk_output( + hidden_states, topk_weights, topk_ids, router_scores + ) + shared = None + with self.stream_fork.scope(enable=get_is_capture_mode()) as fork: + with nvtx_range("moe_experts"): + routed = self.experts( + hidden_states=hidden_states, + topk_output=topk_output, + num_global_tokens=num_global_tokens, + max_num_tokens_per_gpu=max_num_tokens_per_gpu, + ) + if self.routed_scaling_factor != 1.0: + routed *= self.routed_scaling_factor + with fork.branch(): + shared = self._forward_shared_experts(hidden_states) + return routed + shared if shared is not None else routed + + def forward( + self, + hidden_states: torch.Tensor, + input_ids: torch.Tensor, + num_global_tokens: int, + max_num_tokens_per_gpu: int, + ctx: ForwardContext | None = None, + comm_manager: CommManager | None = None, + ) -> torch.Tensor: + if self.use_mega_moe: + return self.forward_mega_moe( + hidden_states, + input_ids, + ctx, + comm_manager, + ) + return self.forward_normal( + hidden_states, input_ids, num_global_tokens, max_num_tokens_per_gpu + ) + + +class DeepseekV4Compressor(nn.Module): + def __init__( + self, + config: PretrainedConfig, + hidden_size: int, + head_dim: int, + compress_ratio: int, + prefix: str, + ) -> None: + super().__init__() + self.compress_ratio = compress_ratio + self.head_dim = head_dim + self.overlap = compress_ratio == 4 + self.coff = 2 if self.overlap else 1 + state_dtype = torch.float32 + self.ape = nn.Parameter( + torch.empty(compress_ratio, self.coff * head_dim, dtype=state_dtype), + requires_grad=False, + ) + self._ape_reordered = False + self.fused_wkv_wgate = MergedColumnParallelLinear( + hidden_size, + [self.coff * head_dim, self.coff * head_dim], + bias=False, + quant_config=None, + prefix=add_prefix("fused_wkv_wgate", prefix), + ) + self.norm = RMSNorm(head_dim, eps=config.rms_norm_eps) + + def process_weights_after_loading(self, module=None) -> None: + del module + if not self.overlap or self._ape_reordered: + return + with torch.no_grad(): + self.ape.data.copy_(_deepseek_v4_reorder_c4_ape_2604(self.ape.data)) + self._ape_reordered = True + + def compute_kv_score(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Pre-compute the input GEMM (fused_wkv_wgate @ hidden_states). + + Can be launched on an aux stream in parallel with the main Q/KV + projection so that ``forward()`` only needs the lightweight state + update + cache write phase. + """ + weight_shape = ( + self.fused_wkv_wgate.output_size_per_partition, + self.fused_wkv_wgate.input_size, + ) + weight = self.fused_wkv_wgate.weight.view(*weight_shape) + if weight.dtype == torch.float8_e4m3fn: + weight = _dequant_fp8_weight(self.fused_wkv_wgate, weight_shape) + kv_score = _deepseek_v4_bf16_linear_fp32(hidden_states, weight) + if kv_score is None: + kv_score = torch.matmul(hidden_states.float(), weight.float().T) + return kv_score + + def forward( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + layer_index: int, + cos_sin_cache: torch.Tensor, + *, + state_cache: torch.Tensor | None = None, + state_block_table: torch.Tensor | None = None, + state_block_size: int | None = None, + state_base_logical_page: torch.Tensor | None = None, + state_slot_mapping: torch.Tensor | None = None, + write_compressed_cache: bool = True, + compressor_slot_cache: dict | None = None, + kv_score: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + pool = ctx.token_to_kv_pool + metadata = _deepseek_v4_forward_metadata(ctx) + if metadata is None: + raise RuntimeError("DeepSeek V4 compressor requires forward metadata") + profile_prefix = ( + f"indexer_compressor_c{self.compress_ratio}" + if not write_compressed_cache + else f"compressor_c{self.compress_ratio}" + ) + if kv_score is None: + with nvtx_range(f"{profile_prefix}_dequant_weight"): + weight_shape = ( + self.fused_wkv_wgate.output_size_per_partition, + self.fused_wkv_wgate.input_size, + ) + weight = self.fused_wkv_wgate.weight.view(*weight_shape) + if weight.dtype == torch.float8_e4m3fn: + weight = _dequant_fp8_weight(self.fused_wkv_wgate, weight_shape) + with nvtx_range(f"{profile_prefix}_matmul"): + kv_score = _deepseek_v4_bf16_linear_fp32(hidden_states, weight) + if kv_score is None: + kv_score = torch.matmul(hidden_states.float(), weight.float().T) + kv, score = kv_score.split([self.coff * self.head_dim] * 2, dim=-1) + if state_cache is None: + state_cache = pool.get_compressor_state_buffer(layer_index) + cache_metadata = metadata.cache + # state/compressed slot mappings depend only on (per-step state, ratio), so reuse + # them across layers of the same ratio within a step. Attn-compressor path only: + # the indexer-compressor passes an explicit state_block_table and is excluded. + memo = compressor_slot_cache if state_block_table is None else None + if state_block_table is None: + state_block_table = cache_metadata.compressor_state_block_tables.get( + self.compress_ratio + ) + state_base_logical_page = ( + cache_metadata.compressor_state_base_logical_pages.get( + self.compress_ratio + ) + ) + if state_block_size is None: + state_block_size = ( + pool.get_compressor_state_block_size(layer_index) + if state_block_table is not None + else pool.state_block_size + ) + valid_token = ( + metadata.is_valid_token[: positions.numel()] + if getattr(metadata, "is_valid_token", None) is not None + else None + ) + if state_slot_mapping is None: + state_hit = ( + memo.get(("state", self.compress_ratio)) if memo is not None else None + ) + if state_hit is not None: + state_slot_mapping, state_block_table = state_hit + else: + if state_block_table is None: + raise RuntimeError( + "DeepSeek V4 missing paged-cache block table for compressor " + f"state ratio={self.compress_ratio}" + ) + state_slot_mapping = _group_slot_mapping_from_raw( + positions, + metadata.token_to_req_indices[: positions.numel()], + state_block_table, + state_block_size, + base_offsets=state_base_logical_page, + ) + state_slot_mapping = _mask_invalid_graph_tokens( + state_slot_mapping, + valid_token, + ) + if memo is not None: + memo[("state", self.compress_ratio)] = ( + state_slot_mapping, + state_block_table, + ) + with nvtx_range(f"{profile_prefix}_save_state"): + save_deepseek_v4_compressor_state( + kv=kv, + score=score, + ape=self.ape, + state_cache=state_cache, + slot_mapping=state_slot_mapping, + positions=positions, + block_size=state_block_size, + compress_ratio=self.compress_ratio, + ) + if not write_compressed_cache: + return kv, score + + kv_cache_block_size = pool.get_compressed_block_size(layer_index) + compressed_hit = ( + memo.get(("compressed", self.compress_ratio)) if memo is not None else None + ) + if compressed_hit is not None: + compressed_slots = compressed_hit + else: + with nvtx_range(f"{profile_prefix}_compressed_slot_mapping"): + compressed_slots = cache_metadata.compressed_slot_mapping( + positions, + self.compress_ratio, + token_to_req_indices=metadata.token_to_req_indices[ + : positions.numel() + ], + query_start_loc=metadata.query_start_loc, + seq_lens=metadata.seq_lens, + kv_cache_block_size=kv_cache_block_size, + use_decode_cache=( + ctx.forward_mode is not None and ctx.forward_mode.is_decode() + ), + is_valid_token=valid_token, + ) + if memo is not None: + memo[("compressed", self.compress_ratio)] = compressed_slots + with nvtx_range(f"{profile_prefix}_cache_insert"): + insert = ( + deepseek_v4_csa_compress_kv_cache_insert + if self.compress_ratio == 4 + else deepseek_v4_hca_compress_kv_cache_insert + ) + insert( + state_cache=state_cache, + token_to_req_indices=metadata.token_to_req_indices[: positions.numel()], + positions=positions, + compressor_slot_mapping=state_slot_mapping, + block_table=state_block_table, + block_table_base_offsets=state_base_logical_page, + compressor_block_size=state_block_size, + rms_norm_weight=self.norm.weight, + rms_norm_eps=self.norm.variance_epsilon, + cos_sin_cache=cos_sin_cache, + kv_cache_2d=pool.get_compressed_kv_buffer_2d(layer_index), + kv_slot_mapping=compressed_slots, + kv_cache_block_size=kv_cache_block_size, + compress_ratio=self.compress_ratio, + ) + return kv, score + + +class DeepseekV4Indexer(nn.Module): + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None, + prefix: str, + compress_ratio: int, + topk_buffer: _DeepseekV4TopKBuffer | None = None, + ) -> None: + super().__init__() + self.wq_b = ReplicatedLinear( + config.q_lora_rank, + config.index_n_heads * config.index_head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("wq_b", prefix), + ) + self.weights_proj = ReplicatedLinear( + config.hidden_size, + config.index_n_heads, + bias=False, + quant_config=None, + prefix=add_prefix("weights_proj", prefix), + ) + self.compressor = DeepseekV4Compressor( + config, + config.hidden_size, + config.index_head_dim, + compress_ratio, + add_prefix("compressor", prefix), + ) + self.use_fp4_cache = _attention_use_fp4_indexer_cache(config) + self.compress_ratio = compress_ratio + self.n_head = int(config.index_n_heads) + self.head_dim = int(config.index_head_dim) + self.topk_tokens = int(config.index_topk) + self.topk_buffer = topk_buffer + self.softmax_scale = self.head_dim**-0.5 + value_bytes = deepseek_v4_indexer_mxfp4_value_bytes(self.head_dim) + scale_bytes = deepseek_v4_indexer_mxfp4_scale_dim(self.head_dim) + self.register_buffer( + "_prefill_gather_values_workspace", + torch.empty((0, value_bytes), dtype=torch.uint8), + persistent=False, + ) + self.register_buffer( + "_prefill_gather_scales_workspace", + torch.empty((0, scale_bytes), dtype=torch.uint8), + persistent=False, + ) + workspace_rows = 1024 * 1024 if self.topk_tokens in (512, 1024, 2048) else 0 + self.register_buffer( + "_persistent_topk_workspace", + torch.empty((workspace_rows,), dtype=torch.uint8), + persistent=False, + ) + + def _prefill_gather_workspace( + self, + rows: int, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor]: + rows = max(0, int(rows)) + value_bytes = deepseek_v4_indexer_mxfp4_value_bytes(self.head_dim) + scale_bytes = deepseek_v4_indexer_mxfp4_scale_dim(self.head_dim) + if ( + self._prefill_gather_values_workspace.device != device + or self._prefill_gather_values_workspace.shape[0] < rows + ): + self._prefill_gather_values_workspace = torch.empty( + (rows, value_bytes), + dtype=torch.uint8, + device=device, + ) + if ( + self._prefill_gather_scales_workspace.device != device + or self._prefill_gather_scales_workspace.shape[0] < rows + ): + self._prefill_gather_scales_workspace = torch.empty( + (rows, scale_bytes), + dtype=torch.uint8, + device=device, + ) + return ( + self._prefill_gather_values_workspace[:rows], + self._prefill_gather_scales_workspace[:rows], + ) + + def prepare_decode_metadata( + self, + *, + positions: torch.Tensor, + metadata: DeepseekV4ForwardMetadata, + ctx: ForwardContext, + indexer_block_size: int, + ) -> None: + if not self.use_fp4_cache or not positions.is_cuda: + return + forward_mode = ctx.forward_mode + if forward_mode is not None and forward_mode.is_mixed(): + num_prefill_tokens = int(metadata.num_prefill_tokens) + num_decode_tokens = metadata.decode_token_count() + elif forward_mode is not None and forward_mode.is_decode(): + num_prefill_tokens = 0 + num_decode_tokens = positions.numel() + else: + return + if num_decode_tokens <= 0: + return + + decode_start = num_prefill_tokens + decode_end = decode_start + num_decode_tokens + decode_positions = positions[decode_start:decode_end] + decode_valid_token = ( + metadata.is_valid_token[decode_start:decode_end] + if getattr(metadata, "is_valid_token", None) is not None + else None + ) + indexer_block_table = metadata.cache.compressed_block_table( + self.compress_ratio, + indexer_block_size, + ) + indexer_block_table_base_offsets = None + if indexer_block_table is not metadata.cache.block_table: + indexer_block_table_base_offsets = ( + metadata.cache.paged_cache_block_table_base_offsets.get( + v4_compressed_kv_group_id(self.compress_ratio) + ) + ) + decode_plan = _deepseek_v4_indexer_decode_plan( + positions=decode_positions, + token_to_req_indices=metadata.token_to_req_indices[decode_start:decode_end], + block_table=indexer_block_table, + cache_block_size=indexer_block_size, + compress_ratio=self.compress_ratio, + metadata=metadata, + is_valid_token=decode_valid_token, + block_table_base_offsets=indexer_block_table_base_offsets, + ) + _deepseek_v4_indexer_decode_schedule_metadata( + positions=decode_positions, + cache_block_size=indexer_block_size, + compress_ratio=self.compress_ratio, + metadata=metadata, + context_lens=decode_plan.context_lens, + ) + + def _forward_sparse_indexer_custom_op( + self, + *, + hidden_states: torch.Tensor, + qr: torch.Tensor, + positions: torch.Tensor, + metadata: DeepseekV4ForwardMetadata, + ctx: ForwardContext, + indexer_cache: torch.Tensor, + indexer_block_size: int, + cos_sin_cache: torch.Tensor, + ) -> torch.Tensor: + if not self.use_fp4_cache: + raise RuntimeError( + "DeepSeek V4 sparse indexer requires MXFP4 indexer cache" + ) + if not positions.is_cuda: + raise RuntimeError("DeepSeek V4 sparse indexer requires CUDA tensors") + + forward_mode = ctx.forward_mode + total_tokens = positions.numel() + if total_tokens == 0: + return torch.empty( + (0, self.topk_tokens), + device=positions.device, + dtype=torch.int32, + ) + num_prefill_tokens, num_decode_tokens = _deepseek_v4_indexer_token_split( + forward_mode, + metadata, + total_tokens, + ) + + with nvtx_range("indexer_wq_b"): + index_q, _ = self.wq_b(qr) + index_q = index_q.view(-1, self.n_head, self.head_dim) + with nvtx_range("indexer_weights_proj"): + weights, _ = self.weights_proj(hidden_states) + with nvtx_range("indexer_prepare_mxfp4"): + packed_index_q, packed_weights = deepseek_v4_prepare_indexer_q_mxfp4( + index_q=index_q, + positions=positions, + cos_sin_cache=cos_sin_cache, + weights=weights, + softmax_scale=self.softmax_scale, + head_scale=self.n_head**-0.5, + ) + + packed_indexer_available = _deepseek_v4_deepgemm_fp4_indexer_available( + packed_index_q[0] + ) + if not packed_indexer_available: + raise RuntimeError( + "DeepSeek V4 sparse indexer requires DeepGEMM FP4 support" + ) + + empty_cpu = torch.empty(0, dtype=torch.int32, device="cpu") + seq_lens_cpu = ( + metadata.seq_lens_cpu[: metadata.num_prefill_reqs] + if metadata.seq_lens_cpu is not None and num_prefill_tokens > 0 + else empty_cpu + ) + query_lens_cpu = ( + metadata.query_lens_cpu[: metadata.num_prefill_reqs] + if metadata.query_lens_cpu is not None and num_prefill_tokens > 0 + else empty_cpu + ) + indexer_block_table = metadata.cache.compressed_block_table( + self.compress_ratio, + indexer_block_size, + ) + prefill_metadata = _deepseek_v4_indexer_prefill_metadata( + metadata=metadata, + block_table=indexer_block_table, + cache_block_size=indexer_block_size, + compress_ratio=self.compress_ratio, + num_prefill_tokens=num_prefill_tokens, + ) + + decode_schedule_metadata = None + decode_plan = None + if num_decode_tokens > 0: + decode_start = num_prefill_tokens + decode_end = decode_start + num_decode_tokens + decode_positions = positions[decode_start:decode_end] + decode_valid_token = ( + metadata.is_valid_token[decode_start:decode_end] + if getattr(metadata, "is_valid_token", None) is not None + else None + ) + decode_plan = _deepseek_v4_indexer_decode_plan( + positions=decode_positions, + token_to_req_indices=metadata.token_to_req_indices[ + decode_start:decode_end + ], + block_table=indexer_block_table, + cache_block_size=indexer_block_size, + compress_ratio=self.compress_ratio, + metadata=metadata, + is_valid_token=decode_valid_token, + ) + decode_schedule_metadata = _deepseek_v4_indexer_decode_schedule_metadata( + positions=decode_positions, + cache_block_size=indexer_block_size, + compress_ratio=self.compress_ratio, + metadata=metadata, + context_lens=decode_plan.context_lens, + ) + indexer_metadata = DeepseekV4SparseIndexerMetadata( + batch_metadata=DeepseekV4IndexerBatchMetadata( + positions=positions, + token_to_req_indices=metadata.token_to_req_indices[:total_tokens], + seq_lens_cpu=seq_lens_cpu, + query_lens_cpu=query_lens_cpu, + num_prefill_tokens=num_prefill_tokens, + num_decode_tokens=num_decode_tokens, + ), + prefill_metadata=prefill_metadata, + decode_plan=decode_plan, + decode_schedule_metadata=decode_schedule_metadata, + ) + prefill_gather_values, prefill_gather_scales = self._prefill_gather_workspace( + prefill_metadata.max_gather_rows(), + positions.device, + ) + + topk_out = ( + self.topk_buffer.get(total_tokens, positions.device) + if self.topk_buffer is not None + else torch.empty( + (total_tokens, self.topk_tokens), + device=positions.device, + dtype=torch.int32, + ) + )[:total_tokens] + return _deepseek_v4_sparse_attn_indexer( + indexer_metadata=indexer_metadata, + indexer_cache=indexer_cache, + indexer_block_table=indexer_block_table, + indexer_block_size=indexer_block_size, + compress_ratio=self.compress_ratio, + packed_q_values=packed_index_q[0], + packed_q_scales=packed_index_q[1], + packed_weights=packed_weights, + topk_indices_buffer=topk_out, + prefill_gather_values_workspace=prefill_gather_values, + prefill_gather_scales_workspace=prefill_gather_scales, + persistent_topk_workspace=self._persistent_topk_workspace, + topk_tokens=self.topk_tokens, + ) + + def forward( + self, + hidden_states: torch.Tensor, + qr: torch.Tensor, + positions: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + layer_index: int, + cos_sin_cache: torch.Tensor, + compressor_slot_cache: dict, + indexer_compressor_kv_score: torch.Tensor | None = None, + ) -> torch.Tensor: + pool = ctx.token_to_kv_pool + metadata = _deepseek_v4_forward_metadata(ctx) + if metadata is None: + raise RuntimeError("DeepSeek V4 indexer requires forward metadata") + cache_metadata = metadata.cache + indexer_state = pool.get_indexer_state_buffer(layer_index) + valid_token = ( + metadata.is_valid_token[: positions.numel()] + if getattr(metadata, "is_valid_token", None) is not None + else None + ) + idx_hit = ( + compressor_slot_cache.get("indexer_state") + if compressor_slot_cache is not None + else None + ) + if idx_hit is not None: + ( + indexer_state_slot_mapping, + indexer_state_block_table, + indexer_state_block_size, + indexer_state_base_logical_page, + ) = idx_hit + else: + indexer_state_block_table = cache_metadata.indexer_state_block_table + indexer_state_base_logical_page = ( + cache_metadata.indexer_state_base_logical_page + ) + if indexer_state_block_table is None: + raise RuntimeError( + "DeepSeek V4 missing paged-cache block table for indexer " + "compressor state" + ) + indexer_state_block_size = pool.get_indexer_state_block_size(layer_index) + indexer_state_slot_mapping = _group_slot_mapping_from_raw( + positions, + metadata.token_to_req_indices[: positions.numel()], + indexer_state_block_table, + indexer_state_block_size, + base_offsets=indexer_state_base_logical_page, + ) + indexer_state_slot_mapping = _mask_invalid_graph_tokens( + indexer_state_slot_mapping, + valid_token, + ) + if compressor_slot_cache is not None: + compressor_slot_cache["indexer_state"] = ( + indexer_state_slot_mapping, + indexer_state_block_table, + indexer_state_block_size, + indexer_state_base_logical_page, + ) + with nvtx_range("indexer_compressor_total"): + self.compressor( + hidden_states=hidden_states, + positions=positions, + ctx=ctx, + out_cache_loc=out_cache_loc, + layer_index=layer_index, + cos_sin_cache=cos_sin_cache, + state_cache=indexer_state, + state_block_table=indexer_state_block_table, + state_block_size=indexer_state_block_size, + state_base_logical_page=indexer_state_base_logical_page, + state_slot_mapping=indexer_state_slot_mapping, + write_compressed_cache=False, + kv_score=indexer_compressor_kv_score, + ) + with nvtx_range("indexer_compressed_slot_mapping"): + indexer_block_size = pool.get_indexer_block_size(layer_index) + compressed_slots = cache_metadata.compressed_slot_mapping( + positions, + self.compress_ratio, + token_to_req_indices=metadata.token_to_req_indices[: positions.numel()], + query_start_loc=metadata.query_start_loc, + seq_lens=metadata.seq_lens, + kv_cache_block_size=indexer_block_size, + use_decode_cache=( + ctx.forward_mode is not None and ctx.forward_mode.is_decode() + ), + is_valid_token=valid_token, + ) + with nvtx_range("indexer_cache_insert"): + deepseek_v4_csa_indexer_cache_insert( + state_cache=indexer_state, + token_to_req_indices=metadata.token_to_req_indices[: positions.numel()], + positions=positions, + compressor_slot_mapping=indexer_state_slot_mapping, + block_table=indexer_state_block_table, + block_table_base_offsets=indexer_state_base_logical_page, + compressor_block_size=indexer_state_block_size, + rms_norm_weight=self.compressor.norm.weight, + rms_norm_eps=self.compressor.norm.variance_epsilon, + cos_sin_cache=cos_sin_cache, + kv_cache_2d=pool.get_indexer_kv_buffer_2d(layer_index), + kv_slot_mapping=compressed_slots, + kv_cache_block_size=indexer_block_size, + use_fp4_cache=self.use_fp4_cache, + compress_ratio=self.compress_ratio, + ) + return self._forward_sparse_indexer_custom_op( + hidden_states=hidden_states, + qr=qr, + positions=positions, + metadata=metadata, + ctx=ctx, + indexer_cache=pool.get_indexer_kv_buffer_2d(layer_index), + indexer_block_size=indexer_block_size, + cos_sin_cache=cos_sin_cache, + ) + + +class DeepseekV4Attention(nn.Module): + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + layer_index: int, + quant_config: QuantizationConfig | None, + prefix: str, + aux_stream: torch.cuda.Stream | None = None, + topk_buffer: _DeepseekV4TopKBuffer | None = None, + cache_layer_index: int | None = None, + ) -> None: + super().__init__() + # `layer_index` addresses checkpoint/config metadata; `cache_layer_index` + # addresses this model's compact KV cache slot. + self.layer_index = layer_index + self.cache_layer_index = ( + layer_index if cache_layer_index is None else cache_layer_index + ) + self.stream_fork = StreamFork(aux_stream) + self.compress_ratio = max(1, int(config.compress_ratios[layer_index])) + if self.compress_ratio <= 1: + self.attention_kind = "swa" + elif self.compress_ratio == 4: + self.attention_kind = "csa" + elif self.compress_ratio == 128: + self.attention_kind = "hca" + else: + raise ValueError( + f"Unsupported DeepSeek V4 compress_ratio={self.compress_ratio}; " + "expected 1, 4, or 128." + ) + self.num_heads = int(config.num_attention_heads) + tp_rank = mapping.attn.tp_rank + tp_size = mapping.attn.tp_size + tp_group = mapping.attn.tp_group + if self.num_heads % tp_size != 0: + raise ValueError( + f"num_attention_heads={self.num_heads} must be divisible " + f"by attn_tp_size={tp_size}" + ) + self.num_local_heads = self.num_heads // tp_size + self.padded_heads = _deepseek_v4_padded_heads(self.num_local_heads) + self.head_dim = int(config.head_dim) + self.qk_rope_head_dim = int(config.qk_rope_head_dim) + self.nope_head_dim = deepseek_v4_nope_dim( + self.head_dim, + self.qk_rope_head_dim, + ) + self.swa_window = int(getattr(config, "sliding_window", 128)) + self.scale = self.head_dim**-0.5 + self.q_lora_rank = config.q_lora_rank + self.o_lora_rank = config.o_lora_rank + self.o_groups = config.o_groups + self.num_local_groups = self.o_groups // tp_size + num_local = self.num_local_heads + self.attn_sink = nn.Parameter( + torch.full((self.padded_heads,), -float("inf"), dtype=torch.float32), + requires_grad=False, + ) + set_weight_attrs( + self.attn_sink, + { + "weight_loader": lambda param, loaded_weight: param.data.__setitem__( + slice(num_local), + loaded_weight[tp_rank * num_local : (tp_rank + 1) * num_local], + ), + }, + ) + rope_base, rope_scaling = deepseek_v4_rope_config(config, self.compress_ratio) + self.rotary_emb = get_rope( + self.qk_rope_head_dim, + rotary_dim=self.qk_rope_head_dim, + max_position=getattr(config, "max_position_embeddings", 8192), + base=rope_base, + rope_scaling=rope_scaling, + is_neox_style=False, + ) + self.indexer_rotary_emb = self.rotary_emb + self.fused_wqa_wkv = MergedColumnParallelLinear( + config.hidden_size, + [self.q_lora_rank, self.head_dim], + bias=False, + quant_config=quant_config, + prefix=add_prefix("fused_wqa_wkv", prefix), + ) + self.q_norm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) + self.wq_b = ColumnParallelLinear( + self.q_lora_rank, + self.num_heads * self.head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("wq_b", prefix), + tp_rank=tp_rank, + tp_size=tp_size, + tp_group=tp_group, + ) + self.kv_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.fused_qkv_norm = FusedRMSNorm(self.q_norm, self.kv_norm) + # Fused QKV-RMSNorm: fuses q_a/kv_a RMSNorm into one kernel instead of two + # sequential RMSNorm + a kv.contiguous() copy. Validated via GSM8K. + self.use_fused_qkv_rmsnorm = True + self.wo_a = ColumnParallelLinear( + self.num_heads * self.head_dim // self.o_groups, + self.o_groups * self.o_lora_rank, + bias=False, + quant_config=quant_config, + prefix=add_prefix("wo_a", prefix), + tp_rank=tp_rank, + tp_size=tp_size, + tp_group=tp_group, + ) + self.wo_a.is_bmm = True + self.wo_a.bmm_batch_size = self.num_local_groups + # SM100 packs the FP8 output-proj scales as INT32 UE8M0 (fp8_einsum + # recipe (1,1,128)); SM90 keeps FP32 block scales (recipe (1,128,128)). + self._o_tma_aligned = torch.cuda.get_device_capability()[0] >= 10 + self.wo_b = RowParallelLinear( + self.o_groups * self.o_lora_rank, + config.hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("wo_b", prefix), + tp_rank=tp_rank, + tp_size=tp_size, + tp_group=tp_group, + ) + if self.compress_ratio > 1: + self.compressor = DeepseekV4Compressor( + config, + config.hidden_size, + self.head_dim, + self.compress_ratio, + add_prefix("compressor", prefix), + ) + else: + self.compressor = None + if self.compress_ratio == 4: + self.indexer = DeepseekV4Indexer( + config, + mapping, + quant_config, + add_prefix("indexer", prefix), + self.compress_ratio, + topk_buffer=topk_buffer, + ) + else: + self.indexer = None + + def _project_q_kv( + self, hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + qr_kv, _ = self.fused_wqa_wkv(hidden_states) + qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1) + if self.use_fused_qkv_rmsnorm and qr.is_cuda and qr.shape[0] > 0: + qr_norm = torch.empty( + qr.shape, + dtype=qr.dtype, + device=qr.device, + ) + kv_norm = torch.empty( + kv.shape, + dtype=kv.dtype, + device=kv.device, + ) + self.fused_qkv_norm( + input_q_a=qr, + input_kv_a=kv, + output_q_a=qr_norm, + output_kv_a=kv_norm, + ) + qr = qr_norm + kv = kv_norm + else: + qr = self.q_norm(qr) + kv = self.kv_norm(kv.contiguous()) + q, _ = self.wq_b(qr) + q = q.view(-1, self.num_local_heads, self.head_dim) + return q, kv, qr + + def _insert_swa_cache( + self, + q: torch.Tensor, + kv: torch.Tensor, + swa_kv_cache: torch.Tensor, + slot_mapping: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + block_size: int, + ) -> None: + # slot_mapping arrives pre-sanitized from _deepseek_v4_swa_slot_mapping + # (invalid tokens and out-of-capacity slots masked to -1); sanitizing + # here again would re-run the mask chain once per layer. + if q.shape[0] == 0: + return + fused_qnorm_rope_kv_insert( + q=q, + kv=kv, + swa_kv_cache_2d=swa_kv_cache.view(swa_kv_cache.shape[0], -1), + slot_mapping=slot_mapping, + positions=positions, + cos_sin_cache=cos_sin_cache, + rms_norm_eps=self.q_norm.variance_epsilon, + block_size=block_size, + ) + + def _project_attention_output( + self, + attn_output: torch.Tensor, + positions: torch.Tensor, + cos_sin_cache: torch.Tensor, + ) -> torch.Tensor: + # Inverse RoPE + grouped block-scaled FP8 quant in one Triton kernel, + # then a single native FP8 GEMM via deep_gemm.fp8_einsum -- replaces the + # inv-rope + per-group online-quant + GEMM chain. wo_a's block scale was + # transformed into the deep_gemm MN-major layout at load time. + heads_per_group = self.num_local_heads // self.num_local_groups + o_fp8, o_scale = deepseek_v4_fused_inv_rope_fp8_quant( + attn_output, + positions, + cos_sin_cache, + n_groups=self.num_local_groups, + heads_per_group=heads_per_group, + nope_dim=self.nope_head_dim, + rope_dim=self.qk_rope_head_dim, + tma_aligned_scales=self._o_tma_aligned, + ) + in_dim = self.num_heads * self.head_dim // self.o_groups + weight = self.wo_a.weight.view(self.num_local_groups, self.o_lora_rank, in_dim) + block_n, block_k = self.wo_a._deep_gemm_block_size + recipe = (1, 1, block_n) if self._o_tma_aligned else (1, block_n, block_k) + z = torch.empty( + (attn_output.shape[0], self.num_local_groups, self.o_lora_rank), + device=attn_output.device, + dtype=torch.bfloat16, + ) + deep_gemm.fp8_einsum( + "bhr,hdr->bhd", + (o_fp8, o_scale), + (weight, self.wo_a.weight_scale_inv), + z, + recipe=recipe, + ) + out, _ = self.wo_b(z.flatten(1)) + return out + + @break_point + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + swa_slot_mapping: torch.Tensor | None = None, + compressor_slot_cache: dict | None = None, + ) -> torch.Tensor: + """DSA attention, one COARSE breakable-graph break point. + + Per layer it does multiple paged-cache writes (SWA / compressor / + indexer), a data/length-dependent indexer -> top-k stage, the FlashMLA + sparse kernel, AND aux-stream forks -- none capturable into a CUDA + graph. Under a prefill-graph capture the whole attention runs eager + (reading the live ``ctx``) while the layer's norms + MoE stay graphed; + direct call otherwise (see ``break_point``). Padding rows are handled + by the existing ``metadata.is_valid_token`` masking, and the + token-shaped inputs are sliced to the real count DSA kernels assert + (see ``slice_to_real_tokens`` below). The cross-layer SWA slot mapping + and compressor memo are shared via ctx -- replay-safe because ctx is + rebound to the live forward (see ``DeepseekV4Model.forward``). + """ + if hidden_states.shape[0] == 0: + return hidden_states + # Cross-layer compressor memo, created once per forward by the first layer. + if compressor_slot_cache is None: + compressor_slot_cache = ctx.dsa_compressor_slot_cache + if compressor_slot_cache is None: + compressor_slot_cache = ctx.dsa_compressor_slot_cache = {} + profile_prefix = f"attn_{self.attention_kind}" + cos_sin_cache = self.rotary_emb.cos_sin_cache + if cos_sin_cache.dtype != torch.float32: + cos_sin_cache = cos_sin_cache.float() + pool = ctx.token_to_kv_pool + metadata = ctx.attn_backend.forward_metadata + if metadata is None: + raise RuntimeError("DeepSeek V4 attention requires forward metadata") + + # Slice padded inputs to the real count the DSA kernels assert (see + # docstring). Only under a breakable capture/replay (ambient ctx set): + # eager forwards are never padded, and the MTP draft path reaches here + # with metadata whose token_to_req_indices does NOT describe q's rows. + token_to_req = getattr(metadata, "token_to_req_indices", None) + if current_forward_ctx() is not None and token_to_req is not None: + positions, hidden_states, out_cache_loc, swa_slot_mapping = ( + slice_to_real_tokens( + token_to_req.numel(), + positions, + hidden_states, + out_cache_loc, + swa_slot_mapping, + ) + ) + + # --- Phase 1: pre-compute input GEMMs in parallel --- + # Q/KV projection on main stream; compressor GEMM(s) on aux stream. + # After this phase, each compressor's forward() receives its + # pre-computed kv_score and skips the internal GEMM, removing the + # shared-weight dependency that prevents safe multi-stream overlap. + compressor_kv_score: torch.Tensor | None = None + indexer_compressor_kv_score: torch.Tensor | None = None + with self.stream_fork.scope( + enable=self.stream_fork.aux_stream is not None + ) as fork: + with nvtx_range(f"{profile_prefix}_project_q_kv"): + q, kv, qr = self._project_q_kv(hidden_states) + with fork.branch(): + if self.compressor is not None: + with nvtx_range(f"{profile_prefix}_compressor_gemm"): + compressor_kv_score = self.compressor.compute_kv_score( + hidden_states + ) + if self.indexer is not None: + with nvtx_range(f"{profile_prefix}_indexer_compressor_gemm"): + indexer_compressor_kv_score = ( + self.indexer.compressor.compute_kv_score(hidden_states) + ) + + if swa_slot_mapping is None: + # Cross-layer SWA slot mapping (per-token, layer-invariant), first layer computes. + swa_slot_mapping = ctx.dsa_swa_slot_mapping + if swa_slot_mapping is None: + swa_slot_mapping = _deepseek_v4_swa_slot_mapping( + ctx, + positions, + out_cache_loc, + ) + ctx.dsa_swa_slot_mapping = swa_slot_mapping + + def insert_swa_cache() -> None: + with nvtx_range(f"{profile_prefix}_insert_swa_cache"): + self._insert_swa_cache( + q=q, + kv=kv, + swa_kv_cache=pool.get_swa_kv_buffer(self.cache_layer_index), + slot_mapping=swa_slot_mapping, + positions=positions, + cos_sin_cache=cos_sin_cache, + block_size=pool.swa_block_size, + ) + + def run_compressor() -> None: + if self.compressor is None: + return + with nvtx_range(f"{profile_prefix}_compressor"): + self.compressor( + hidden_states=hidden_states, + positions=positions, + ctx=ctx, + out_cache_loc=out_cache_loc, + layer_index=self.cache_layer_index, + cos_sin_cache=cos_sin_cache, + compressor_slot_cache=compressor_slot_cache, + kv_score=compressor_kv_score, + ) + + # --- Phase 2: state-update + cache-write overlap --- + # With GEMMs already done, the remaining work per stream is lightweight + # state updates and paged cache writes — safe to overlap. + topk_indices = None + if self.indexer is not None: + if self.compressor is None: + raise RuntimeError("compressor is required when indexer is enabled.") + with nvtx_range(f"{profile_prefix}_indexer_prepare_decode_metadata"): + self.indexer.prepare_decode_metadata( + positions=positions, + metadata=metadata, + ctx=ctx, + indexer_block_size=pool.get_indexer_block_size( + self.cache_layer_index + ), + ) + + with self.stream_fork.scope( + enable=self.stream_fork.aux_stream is not None + ) as fork: + with nvtx_range(f"{profile_prefix}_indexer"): + topk_indices = self.indexer( + hidden_states=hidden_states, + qr=qr, + positions=positions, + ctx=ctx, + out_cache_loc=out_cache_loc, + layer_index=self.cache_layer_index, + cos_sin_cache=cos_sin_cache, + compressor_slot_cache=compressor_slot_cache, + indexer_compressor_kv_score=indexer_compressor_kv_score, + ) + with fork.branch(): + insert_swa_cache() + run_compressor() + elif self.compressor is not None: + with self.stream_fork.scope( + enable=self.stream_fork.aux_stream is not None + ) as fork: + run_compressor() + with fork.branch(): + insert_swa_cache() + else: + insert_swa_cache() + forward_mode = ctx.forward_mode + if forward_mode is None: + raise RuntimeError("DeepSeek V4 attention requires forward mode") + if forward_mode.is_mixed(): + with nvtx_range(f"{profile_prefix}_mixed_backend"): + attn_output = ctx.attn_backend.forward_deepseek_v4_mixed( + q=q, + positions=positions, + token_to_kv_pool=pool, + layer_id=self.cache_layer_index, + kind=self.attention_kind, + compress_ratio=self.compress_ratio, + num_local_heads=self.num_local_heads, + padded_heads=self.padded_heads, + head_dim=self.head_dim, + window_size=self.swa_window, + softmax_scale=self.scale, + attn_sink=self.attn_sink, + topk_indices=topk_indices, + ) + elif forward_mode.is_decode(): + with nvtx_range(f"{profile_prefix}_decode_backend"): + attn_output = ctx.attn_backend.forward_deepseek_v4_decode( + q=q, + positions=positions, + token_to_kv_pool=pool, + layer_id=self.cache_layer_index, + kind=self.attention_kind, + compress_ratio=self.compress_ratio, + num_local_heads=self.num_local_heads, + padded_heads=self.padded_heads, + head_dim=self.head_dim, + window_size=self.swa_window, + softmax_scale=self.scale, + attn_sink=self.attn_sink, + topk_indices=topk_indices, + ) + elif forward_mode.is_extend_or_mixed(): + with nvtx_range(f"{profile_prefix}_prefill_backend"): + attn_output = ctx.attn_backend.forward_deepseek_v4_prefill( + q=q, + positions=positions, + token_to_kv_pool=pool, + layer_id=self.cache_layer_index, + kind=self.attention_kind, + compress_ratio=self.compress_ratio, + num_local_heads=self.num_local_heads, + padded_heads=self.padded_heads, + head_dim=self.head_dim, + window_size=self.swa_window, + softmax_scale=self.scale, + attn_sink=self.attn_sink, + topk_indices=topk_indices, + ) + else: + raise RuntimeError(f"Unsupported DeepSeek V4 forward mode: {forward_mode}") + with nvtx_range(f"{profile_prefix}_output_proj"): + return self._project_attention_output( + attn_output, + positions, + cos_sin_cache, + ) + + +class DeepseekV4DecoderLayer(nn.Module): + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + mapping: Mapping, + quant_config: QuantizationConfig | None, + prefix: str, + aux_stream: torch.cuda.Stream | None = None, + topk_buffer: _DeepseekV4TopKBuffer | None = None, + cache_layer_index: int | None = None, + ) -> None: + super().__init__() + self.mapping = mapping + self.layer_id = layer_id + self.hidden_size = config.hidden_size + self.rms_norm_eps = config.rms_norm_eps + self.hc_mult = config.hc_mult + self.hc_sinkhorn_iters = config.hc_sinkhorn_iters + self.hc_eps = config.hc_eps + mix_hc = (2 + self.hc_mult) * self.hc_mult + hc_dim = self.hc_mult * config.hidden_size + self.attn = DeepseekV4Attention( + config, + mapping, + layer_id, + quant_config, + add_prefix("attn", prefix), + aux_stream=aux_stream, + topk_buffer=topk_buffer, + cache_layer_index=cache_layer_index, + ) + self.ffn = DeepseekV4MoE( + config, + mapping, + quant_config, + layer_id, + add_prefix("ffn", prefix), + aux_stream=aux_stream, + ) + self.comm_manager = CommManager( + mapping=mapping, + layer_id=layer_id, + is_moe=True, + prev_is_moe=True, + ) + self.attn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.ffn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hc_attn_fn = nn.Parameter( + torch.empty(mix_hc, hc_dim, dtype=torch.float32), requires_grad=False + ) + self.hc_ffn_fn = nn.Parameter( + torch.empty(mix_hc, hc_dim, dtype=torch.float32), requires_grad=False + ) + self.hc_attn_base = nn.Parameter( + torch.empty(mix_hc, dtype=torch.float32), requires_grad=False + ) + self.hc_ffn_base = nn.Parameter( + torch.empty(mix_hc, dtype=torch.float32), requires_grad=False + ) + self.hc_attn_scale = nn.Parameter( + torch.empty(3, dtype=torch.float32), requires_grad=False + ) + self.hc_ffn_scale = nn.Parameter( + torch.empty(3, dtype=torch.float32), requires_grad=False + ) + + def _pre_mlp_input_ids_comm( + self, input_ids: torch.Tensor, ctx: ForwardContext + ) -> torch.Tensor: + if not self.mapping.moe.has_tp_ep: + return input_ids + if self.comm_manager.use_all_reduce(is_moe=True): + return input_ids + + token_counts = self.comm_manager.moe_tp_ep_group_scattered_num_tokens(ctx) + max_tokens = max(token_counts) + padded = torch.empty( + (max_tokens,), device=input_ids.device, dtype=input_ids.dtype + ) + padded[: input_ids.shape[0]].copy_(input_ids) + if input_ids.shape[0] < max_tokens: + padded[input_ids.shape[0] :].zero_() + + gathered = [torch.empty_like(padded) for _ in token_counts] + group = pg_manager.get_process_group("nccl", self.mapping.moe.tp_ep_group) + torch.distributed.all_gather(gathered, padded, group=group) + return torch.cat( + [tokens[:count] for tokens, count in zip(gathered, token_counts)], dim=0 + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + input_ids: torch.Tensor, + swa_slot_mapping: torch.Tensor | None = None, + compressor_slot_cache: dict | None = None, + hc_x_prev: torch.Tensor | None = None, + hc_post_prev: torch.Tensor | None = None, + hc_comb_prev: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if hc_x_prev is not None: + with nvtx_range("hc_fused_attn_pre"): + residual, hidden_states, post, comb = mhc_fused_hc( + hc_x_prev, + hidden_states, + hc_post_prev, + hc_comb_prev, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_sinkhorn_iters, + ) + else: + residual = hidden_states + with nvtx_range("hc_attn_pre"): + hidden_states, post, comb = mhc_pre( + hidden_states, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_sinkhorn_iters, + ) + with nvtx_range("attn_norm"): + hidden_states = self.attn_norm(hidden_states) + + with nvtx_range("attn_total"): + hidden_states = self.attn( + positions, + hidden_states, + ctx, + out_cache_loc, + swa_slot_mapping, + compressor_slot_cache, + ) + + with nvtx_range("hc_fused_ffn_pre"): + residual, hidden_states, post, comb = mhc_fused_hc( + hidden_states, + residual, + post, + comb, + self.hc_ffn_fn, + self.hc_ffn_scale, + self.hc_ffn_base, + self.rms_norm_eps, + self.hc_eps, + self.hc_sinkhorn_iters, + ) + with nvtx_range("ffn_norm"): + hidden_states = self.ffn_norm(hidden_states) + + ffn_input_ids = input_ids + use_mega_moe = getattr(self.ffn, "use_mega_moe", False) + if use_mega_moe: + token_counts = self.comm_manager.moe_tp_ep_group_scattered_num_tokens(ctx) + num_global_tokens = sum(token_counts) + max_num_tokens_per_gpu = max(token_counts) if token_counts else 0 + else: + token_counts = None + with nvtx_range("pre_mlp_comm"): + hidden_states = self.comm_manager.pre_mlp_comm(hidden_states, ctx) + if self.ffn.gate.is_hash_moe: + with nvtx_range("pre_mlp_input_ids_comm"): + ffn_input_ids = self._pre_mlp_input_ids_comm(input_ids, ctx) + with nvtx_range("moe_get_num_tokens"): + num_global_tokens, max_num_tokens_per_gpu = ( + self.comm_manager.get_num_tokens(ctx) + ) + with nvtx_range("ffn_total"): + hidden_states = self.ffn( + hidden_states, + ffn_input_ids, + num_global_tokens, + max_num_tokens_per_gpu, + ctx=ctx if use_mega_moe else None, + comm_manager=self.comm_manager if use_mega_moe else None, + ) + if not use_mega_moe: + with nvtx_range("post_mlp_comm"): + hidden_states, _ = self.comm_manager.post_mlp_comm( + hidden_states, None, ctx + ) + # Defer ffn post_mapping to next layer's fused_hc + return residual, hidden_states, post, comb + + +class DeepseekV4Model(nn.Module): + fall_back_to_pt_during_load = False + + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.mapping = mapping + self.hc_mult = config.hc_mult + self.hc_eps = config.hc_eps + self.rms_norm_eps = config.rms_norm_eps + self.aux_stream = torch.cuda.Stream() if torch.cuda.is_available() else None + self.topk_indices_buffer = _DeepseekV4TopKBuffer(int(config.index_topk)) + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + tp_rank=mapping.attn.tp_rank, + tp_size=mapping.attn.tp_size, + tp_group=mapping.attn.tp_group, + prefix=add_prefix("embed_tokens", prefix), + ) + self.layers = nn.ModuleList( + [ + DeepseekV4DecoderLayer( + config, + layer_id, + mapping, + quant_config, + add_prefix(f"layers.{layer_id}", prefix), + aux_stream=self.aux_stream, + topk_buffer=self.topk_indices_buffer, + ) + for layer_id in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + hc_dim = config.hc_mult * config.hidden_size + self.hc_head_fn = nn.Parameter( + torch.empty(config.hc_mult, hc_dim, dtype=torch.float32), + requires_grad=False, + ) + self.hc_head_base = nn.Parameter( + torch.empty(config.hc_mult, dtype=torch.float32), requires_grad=False + ) + self.hc_head_scale = nn.Parameter( + torch.empty(1, dtype=torch.float32), requires_grad=False + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + input_embeds: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, list[torch.Tensor] | None]: + hidden_states = input_embeds + if hidden_states is None: + with nvtx_range("embed_tokens"): + hidden_states = self.embed_tokens(input_ids) + with nvtx_range("hc_repeat"): + hidden_states = hidden_states.unsqueeze(1).repeat(1, self.hc_mult, 1) + # Layer-invariant DSA memos: computing them HERE (graphed scope) would bake + # dummy addresses; the first attention break computes them LIVE onto ctx. + ctx.dsa_swa_slot_mapping = None + ctx.dsa_compressor_slot_cache = None + hc_x_prev = None + hc_post_prev = None + hc_comb_prev = None + for layer in self.layers: + hidden_states, hc_x_prev, hc_post_prev, hc_comb_prev = layer( + positions, + hidden_states, + ctx, + out_cache_loc, + input_ids, + None, # swa_slot_mapping: first break computes + caches on ctx + None, # compressor_slot_cache: shared dict created on ctx + hc_x_prev, + hc_post_prev, + hc_comb_prev, + ) + with nvtx_range("hc_ffn_post_final"): + hidden_states = mhc_post( + hc_x_prev, hidden_states, hc_post_prev, hc_comb_prev + ) + aux_hidden_states = None + if ( + ctx.capture_hidden_mode is not None + and ctx.capture_hidden_mode.need_capture() + ): + # V4 MTP consumes the pre-hc_head hypercompressed residual. + aux_hidden_states = [hidden_states.flatten(1)] + with nvtx_range("hc_head"): + hidden_states = hc_head( + hidden_states, + self.hc_head_fn, + self.hc_head_scale, + self.hc_head_base, + self.rms_norm_eps, + self.hc_eps, + ) + with nvtx_range("final_norm"): + hidden_states = self.norm(hidden_states) + return hidden_states, aux_hidden_states + + +class DeepseekV4ForCausalLM(BaseCausalLM): + model_cls = DeepseekV4Model + + def get_stacked_params_mapping(self): + return [ + ("gate_up_proj", "w1", 0), + ("gate_up_proj", "w3", 1), + ("attn.fused_wqa_wkv", "attn.wq_a", 0), + ("attn.fused_wqa_wkv", "attn.wkv", 1), + ("compressor.fused_wkv_wgate", "compressor.wkv", 0), + ("compressor.fused_wkv_wgate", "compressor.wgate", 1), + ] + + @staticmethod + def _map_weight_name(name: str) -> str: + if name.startswith("layers."): + name = "model." + name + elif name.startswith("embed."): + name = name.replace("embed.", "model.embed_tokens.", 1) + elif name.startswith("norm."): + name = "model." + name + elif name.startswith("hc_head"): + name = "model." + name + elif name == "head.weight": + name = "lm_head.weight" + if ".shared_experts.w2" in name: + name = name.replace(".shared_experts.w2", ".shared_experts.down_proj") + if ".ffn.gate.bias" in name: + name = name.replace(".ffn.gate.bias", ".ffn.gate.e_score_correction_bias") + if re.search(r"\.experts\.\d+\.w[123]\.scale$", name): + name = name.replace(".scale", ".weight_scale") + elif name.endswith(".scale"): + name = name[:-6] + ".weight_scale_inv" + return name + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + stacked_params_mapping = self.get_stacked_params_mapping() + params_dict = dict(self.named_parameters()) + moe_loader = build_moe_checkpoint_loader( + params_dict=params_dict, + expert_schema=ExpertCheckpointSchema( + gate_proj_name="w1", + down_proj_name="w2", + up_proj_name="w3", + ), + num_experts=self.config.n_routed_experts, + ep_rank=self.mapping.moe.ep_rank, + ep_size=self.mapping.moe.ep_size, + ) + for raw_name, loaded_weight in weights: + name = self._map_weight_name(raw_name) + if name.startswith("mtp."): + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name or ".experts." in name: + continue + name = name.replace(weight_name, param_name) + param = params_dict.get(name) + if param is None: + break + param.weight_loader(param, loaded_weight, shard_id) + break + else: + if moe_loader.matches(name): + moe_loader.load(name, loaded_weight) + continue + param = params_dict.get(name) + if param is None: + continue + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + del params_dict, moe_loader + self.post_load_weights() + self.warmup_deep_gemm() + + def post_load_weights(self): + mega_moe_experts: list[DeepseekV4MegaMoEExperts] = [] + for module in self.modules(): + if isinstance(module, DeepseekV4Compressor): + module.process_weights_after_loading() + elif isinstance(module, DeepseekV4MegaMoEExperts): + module.finalize_weights() + mega_moe_experts.append(module) + elif isinstance(module, MoELayer): + module.process_weights_after_loading(module) + + def warmup_deep_gemm(self) -> None: + """Pre-compile all DeepGEMM JIT kernels used by this model. + + Called after post_load_weights (not inside it) so that + finalize_weights temporaries have been freed by GC and there is + enough GPU memory for the symmetric buffer allocation. + """ + if deep_gemm is None: + return + import gc + + gc.collect() + torch.cuda.empty_cache() + + self._warmup_mega_moe_jit() + self._warmup_prefill_jit() + + def _warmup_mega_moe_jit(self) -> None: + if os.environ.get("TOKENSPEED_DISABLE_MEGA_MOE_WARMUP") == "1": + return + for module in self.modules(): + if not isinstance(module, DeepseekV4MegaMoEExperts): + continue + if module._transformed_l1_weights is None: + continue + logger.info("Pre-compiling DeepGEMM mega-MoE kernel variants...") + group = pg_manager.get_process_group( + "nccl", + module.mapping.moe.tp_ep_group, + ) + if torch.distributed.is_initialized(): + torch.distributed.barrier(group=group) + deep_gemm.warmup_mega_moe_jit( + num_experts=module.num_experts, + max_num_tokens=module.max_num_tokens, + top_k=module.top_k, + hidden_size=module.hidden_size, + device=torch.device("cuda", torch.cuda.current_device()), + transformed_l1_weights=module._transformed_l1_weights, + transformed_l2_weights=module._transformed_l2_weights, + symm_buffer=module.get_symm_buffer(), + activation_clamp=module.swiglu_limit, + ) + return + + def _warmup_prefill_jit(self) -> None: + if deep_gemm is None: + return + if torch.cuda.get_device_capability()[0] < 10: + return + config = self.config + tp_size = self.mapping.attn.tp_size if self.mapping else 1 + logger.info("Pre-compiling DeepGEMM prefill kernel variants...") + deep_gemm.warmup_prefill_jit( + hidden_size=config.hidden_size, + num_attention_heads=config.num_attention_heads, + head_dim=getattr(config, "head_dim", 128), + hc_mult=getattr(config, "hc_mult", 0), + kv_lora_rank=getattr(config, "kv_lora_rank", 0), + index_n_heads=getattr(config, "index_n_heads", 0), + index_head_dim=getattr(config, "index_head_dim", 0), + indexer_cache_block_size=V4_KERNEL_BLOCK_ROWS, + # With MTP/speculation the verify step flattens bs*num_draft_tokens + # into the decode indexer's num_tokens, so the paged-logits metadata + # kernel (JIT-keyed on the 32-aligned token count) hits larger + # buckets than plain decode. Scale the warmup ceiling by the draft + # factor so those buckets are covered (no-op when speculation is off). + max_decode_tokens=max( + int(global_server_args_dict.get("max_cudagraph_capture_size", 0) or 0), + int(global_server_args_dict.get("max_num_seqs", 0) or 0), + 1, + ) + * ( + int(global_server_args_dict.get("speculative_num_draft_tokens", 1) or 1) + if global_server_args_dict.get("speculative_algorithm") + else 1 + ), + mxfp4_block_size=DEEPSEEK_V4_MXFP4_BLOCK_SIZE, + tp_size=tp_size, + # Prefill GEMM/prenorm M is capped per forward by chunked_prefill_size + # (continuous batching), the same ceiling mega_moe warms to. Hardcoding + # 8192 would leave M in (8192, chunked_prefill_size] to JIT inline. + max_tokens=_deepseek_v4_mega_moe_max_num_tokens(), + device=torch.device("cuda", torch.cuda.current_device()), + ) + + def post_quant_warmup(self) -> None: + """Called by the weight loader after all quant process_weights_after_loading.""" + if deep_gemm is not None: + deep_gemm.warmup_fp8_gemm_nt_from_model( + self, max_tokens=_deepseek_v4_mega_moe_max_num_tokens() + ) + + @classmethod + def get_model_config_for_expert_location(cls, config): + return ModelConfigForExpertLocation( + num_layers=config.num_hidden_layers, + num_logical_experts=config.n_routed_experts, + num_groups=getattr(config, "n_group", 0), + ) + + +EntryClass = [ + DeepseekV4ForCausalLM, +] diff --git a/python/tokenspeed/runtime/models/deepseek_v4_mtp.py b/python/tokenspeed/runtime/models/deepseek_v4_mtp.py new file mode 100644 index 0000000..7376b85 --- /dev/null +++ b/python/tokenspeed/runtime/models/deepseek_v4_mtp.py @@ -0,0 +1,524 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Inference-only DeepSeek V4 MTP / NextN draft model.""" + +from __future__ import annotations + +import logging +import re +from collections.abc import Iterable + +import torch +from torch import nn +from transformers import PretrainedConfig + +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.layers.layernorm import RMSNorm +from tokenspeed.runtime.layers.linear import ReplicatedLinear +from tokenspeed.runtime.layers.logits_processor import LogitsMetadata, LogitsProcessor +from tokenspeed.runtime.layers.moe import ( + ExpertCheckpointSchema, + build_moe_checkpoint_loader, +) +from tokenspeed.runtime.layers.moe.expert import MoELayer +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader +from tokenspeed.runtime.models.deepseek_v4 import ( + DeepseekV4Compressor, + DeepseekV4DecoderLayer, + DeepseekV4MegaMoEExperts, + _deepseek_v4_swa_slot_mapping, + hc_head, + mhc_post, +) +from tokenspeed.runtime.utils import add_prefix + +logger = logging.getLogger(__name__) + + +_EXPERT_SCALE_RE = re.compile(r"\.experts\.\d+\.w[123]\.scale$") + + +def _spec_layer_idx(config: PretrainedConfig, weight_name: str) -> int | None: + if getattr(config, "num_nextn_predict_layers", 0) <= 0: + return None + start = config.num_hidden_layers + for idx in range(start, start + config.num_nextn_predict_layers): + if weight_name.startswith(f"model.layers.{idx}."): + return idx + return None + + +def _find_mtp_layer_idx(name: str) -> int: + parts = name.split(".") + if len(parts) > 1 and parts[0] == "mtp": + try: + return int(parts[1]) + except ValueError: + pass + for part in parts: + try: + return int(part) + except ValueError: + continue + return 0 + + +class DeepseekV4MTPSharedHead(nn.Module): + def __init__(self, config: PretrainedConfig) -> None: + super().__init__() + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + +class DeepseekV4MultiTokenPredictorLayer(nn.Module): + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + layer_id: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + cache_layer_index: int | None = None, + ) -> None: + super().__init__() + self.config = config + self.layer_id = layer_id + self.cache_layer_index = ( + layer_id if cache_layer_index is None else cache_layer_index + ) + self.rms_norm_eps = config.rms_norm_eps + self.hc_eps = config.hc_eps + self.hc_mult = config.hc_mult + + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.e_proj = ReplicatedLinear( + config.hidden_size, + config.hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("e_proj", prefix), + ) + self.h_proj = ReplicatedLinear( + config.hidden_size, + config.hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("h_proj", prefix), + ) + self.hc_head_fn = nn.Parameter( + torch.empty( + self.hc_mult, + self.hc_mult * config.hidden_size, + dtype=torch.float32, + ), + requires_grad=False, + ) + self.hc_head_base = nn.Parameter( + torch.empty(self.hc_mult, dtype=torch.float32), + requires_grad=False, + ) + self.hc_head_scale = nn.Parameter( + torch.empty(1, dtype=torch.float32), + requires_grad=False, + ) + self.shared_head = DeepseekV4MTPSharedHead(config) + self.mtp_block = DeepseekV4DecoderLayer( + config, + layer_id, + mapping, + quant_config, + add_prefix("mtp_block", prefix), + cache_layer_index=self.cache_layer_index, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + input_embeds: torch.Tensor | None = None, + ) -> torch.Tensor: + if input_embeds is None: + raise ValueError("DeepSeek V4 MTP requires input_embeds.") + input_embeds = torch.where(positions.unsqueeze(-1) == 0, 0, input_embeds) + input_embeds = self.enorm(input_embeds) + previous_hidden_states = previous_hidden_states.view( + -1, self.hc_mult, self.config.hidden_size + ) + previous_hidden_states = self.hnorm(previous_hidden_states) + h_out, _ = self.h_proj(previous_hidden_states) + e_out, _ = self.e_proj(input_embeds) + hidden_states = h_out + e_out.unsqueeze(-2) + + swa_slot_mapping = _deepseek_v4_swa_slot_mapping( + ctx, + positions, + out_cache_loc, + ) + residual, x_def, post_def, comb_def = self.mtp_block( + positions, + hidden_states, + ctx, + out_cache_loc, + input_ids, + swa_slot_mapping, + ) + return mhc_post(x_def, residual, post_def, comb_def) + + def compute_logits_hidden(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = hidden_states.view(-1, self.hc_mult, self.config.hidden_size) + hidden_states = hc_head( + hidden_states, + self.hc_head_fn, + self.hc_head_scale, + self.hc_head_base, + self.rms_norm_eps, + self.hc_eps, + ) + return self.shared_head.norm(hidden_states) + + +class DeepseekV4MultiTokenPredictor(nn.Module): + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.mapping = mapping + self.mtp_start_layer_idx = config.num_hidden_layers + self.num_mtp_layers = config.num_nextn_predict_layers + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + tp_rank=mapping.attn.tp_rank, + tp_size=mapping.attn.tp_size, + tp_group=mapping.attn.tp_group, + prefix=add_prefix("embed_tokens", prefix), + ) + layers = {} + for local_idx in range(self.num_mtp_layers): + # Checkpoint layer ids remain global, while draft KV slots are compact. + layer_idx = self.mtp_start_layer_idx + local_idx + layers[str(layer_idx)] = DeepseekV4MultiTokenPredictorLayer( + config, + mapping, + layer_idx, + quant_config, + add_prefix(f"layers.{layer_idx}", prefix), + cache_layer_index=local_idx, + ) + self.layers = nn.ModuleDict(layers) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + input_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + if input_embeds is None: + input_embeds = self.embed_tokens(input_ids) + current_step_idx = spec_step_idx % self.num_mtp_layers + layer_idx = self.mtp_start_layer_idx + current_step_idx + return self.layers[str(layer_idx)]( + input_ids, + positions, + previous_hidden_states, + ctx, + out_cache_loc, + input_embeds, + ) + + def compute_logits_hidden( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor: + current_step_idx = spec_step_idx % self.num_mtp_layers + layer_idx = self.mtp_start_layer_idx + current_step_idx + return self.layers[str(layer_idx)].compute_logits_hidden(hidden_states) + + +class DeepseekV4ForCausalLMNextN(nn.Module): + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + nn.Module.__init__(self) + self.config = config + self.mapping = mapping + self.quant_config = quant_config + self.model = DeepseekV4MultiTokenPredictor( + config, + mapping=mapping, + quant_config=quant_config, + prefix=add_prefix("model", prefix), + ) + if self.mapping.attn.has_dp: + self.lm_head = ReplicatedLinear( + config.hidden_size, + config.vocab_size, + bias=False, + prefix=add_prefix("lm_head", prefix), + ) + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + prefix=add_prefix("lm_head", prefix), + ) + self.logits_processor = LogitsProcessor( + config, + skip_all_gather=self.mapping.attn.has_dp, + do_argmax=True, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + + def get_hot_token_id(self): + return None + + def get_embed_and_head(self) -> tuple[torch.Tensor, torch.Tensor]: + return self.model.embed_tokens.weight, self.lm_head.weight + + def set_embed_and_head(self, embed: torch.Tensor, head: torch.Tensor) -> None: + del self.model.embed_tokens.weight + del self.lm_head.weight + self.model.embed_tokens.weight = embed + self.lm_head.weight = head + torch.cuda.empty_cache() + torch.cuda.synchronize() + + @torch.no_grad() + def forward( + self, + ctx: ForwardContext, + input_ids: torch.Tensor, + positions: torch.Tensor, + out_cache_loc: torch.Tensor, + input_embeds: torch.Tensor | None = None, + captured_hidden_states: torch.Tensor | None = None, + spec_step_idx: int = 0, + **kwargs, + ): + del kwargs + if captured_hidden_states is None: + if not ctx.forward_mode.is_idle(): + raise ValueError("DeepSeek V4 MTP requires captured_hidden_states.") + captured_hidden_states = torch.zeros( + 0, + self.config.hc_mult * self.config.hidden_size, + device=input_ids.device, + dtype=self.model.embed_tokens.weight.dtype, + ) + + mtp_hidden_states = self.model( + input_ids, + positions, + captured_hidden_states, + ctx, + out_cache_loc, + input_embeds=input_embeds, + spec_step_idx=spec_step_idx, + ).flatten(1) + logits_hidden_states = self.model.compute_logits_hidden( + mtp_hidden_states, + spec_step_idx, + ) + logits_metadata = LogitsMetadata.from_forward_context(ctx) + return self.logits_processor( + input_ids, + logits_hidden_states, + self.lm_head, + logits_metadata, + aux_hidden_states=[mtp_hidden_states], + ) + + @staticmethod + def _remap_weight_name(name: str) -> str: + for old, new in { + ".emb.tok_emb.weight": ".embed_tokens.weight", + ".head.weight": ".shared_head.head.weight", + ".norm.weight": ".shared_head.norm.weight", + }.items(): + if old in name: + name = name.replace(old, new) + return name + + @staticmethod + def _rewrite_spec_layer_name(spec_layer: int, name: str) -> str: + spec_layer_weight_names = ( + "embed_tokens", + "enorm", + "hnorm", + "h_proj", + "e_proj", + "shared_head", + "hc_head_fn", + "hc_head_base", + "hc_head_scale", + ) + shared_weight_names = ("embed_tokens",) + is_spec_weight = any( + weight_name in name for weight_name in spec_layer_weight_names + ) + is_shared_weight = any( + weight_name in name for weight_name in shared_weight_names + ) + if not is_spec_weight: + name = name.replace( + f"model.layers.{spec_layer}.", + f"model.layers.{spec_layer}.mtp_block.", + ) + elif is_shared_weight: + name = name.replace(f"model.layers.{spec_layer}.", "model.") + return name + + def _map_checkpoint_name(self, raw_name: str) -> str | None: + if raw_name.startswith("mtp."): + mtp_layer_idx = _find_mtp_layer_idx(raw_name) + raw_name = raw_name.replace( + f"mtp.{mtp_layer_idx}.", + f"model.layers.{self.config.num_hidden_layers + mtp_layer_idx}.", + 1, + ) + spec_layer = _spec_layer_idx(self.config, raw_name) + if spec_layer is None: + return None + name = self._remap_weight_name(raw_name) + name = self._rewrite_spec_layer_name(spec_layer, name) + if name.endswith(".shared_head.head.weight"): + return None + if name.endswith(".scale"): + suffix = ( + ".weight_scale" + if _EXPERT_SCALE_RE.search(name) + else ".weight_scale_inv" + ) + name = name.removesuffix(".scale") + suffix + if ".shared_experts.w2" in name: + name = name.replace(".shared_experts.w2", ".shared_experts.down_proj") + if ".ffn.gate.bias" in name: + name = name.replace(".ffn.gate.bias", ".ffn.gate.e_score_correction_bias") + return name + + def get_stacked_params_mapping(self): + return [ + ("gate_up_proj", "w1", 0), + ("gate_up_proj", "w3", 1), + ("attn.fused_wqa_wkv", "attn.wq_a", 0), + ("attn.fused_wqa_wkv", "attn.wkv", 1), + ("compressor.fused_wkv_wgate", "compressor.wkv", 0), + ("compressor.fused_wkv_wgate", "compressor.wgate", 1), + ] + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + stacked_params_mapping = self.get_stacked_params_mapping() + params_dict = dict(self.named_parameters()) + moe_loader = build_moe_checkpoint_loader( + params_dict=params_dict, + expert_schema=ExpertCheckpointSchema( + gate_proj_name="w1", + down_proj_name="w2", + up_proj_name="w3", + ), + num_experts=self.config.n_routed_experts, + ep_rank=self.mapping.moe.ep_rank, + ep_size=self.mapping.moe.ep_size, + ) + loaded_params: set[str] = set() + for raw_name, loaded_weight in weights: + name = self._map_checkpoint_name(raw_name) + if name is None: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name or ".experts." in name: + continue + mapped_name = name.replace(weight_name, param_name) + param = params_dict.get(mapped_name) + if param is None: + break + param.weight_loader(param, loaded_weight, shard_id) + loaded_params.add(mapped_name) + break + else: + if moe_loader.matches(name): + mapped_name = moe_loader.load(name, loaded_weight) + loaded_params.add(mapped_name) + continue + param = params_dict.get(name) + if param is None: + logger.debug("Skipping unmatched DeepSeek V4 MTP weight: %s", name) + continue + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + missing_layers = [] + for layer_idx in range( + self.model.mtp_start_layer_idx, + self.model.mtp_start_layer_idx + self.model.num_mtp_layers, + ): + if not any(f"model.layers.{layer_idx}." in name for name in loaded_params): + missing_layers.append(layer_idx) + if missing_layers: + raise ValueError( + "DeepSeek V4 MTP weights missing for speculative layer(s) " + f"{missing_layers}. Use a checkpoint that includes `mtp.*` " + "weights or disable NEXTN speculative decoding." + ) + self.post_load_weights() + return loaded_params + + def post_load_weights(self): + for module in self.modules(): + if isinstance(module, DeepseekV4Compressor): + module.process_weights_after_loading() + elif isinstance(module, DeepseekV4MegaMoEExperts): + module.finalize_weights() + elif isinstance(module, MoELayer): + module.process_weights_after_loading(module) + + +EntryClass = [DeepseekV4ForCausalLMNextN] diff --git a/python/tokenspeed/runtime/models/dflash.py b/python/tokenspeed/runtime/models/dflash.py new file mode 100644 index 0000000..e68a25d --- /dev/null +++ b/python/tokenspeed/runtime/models/dflash.py @@ -0,0 +1,430 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from collections.abc import Iterable + +import torch +from torch import nn + +from tokenspeed.runtime.distributed.comm_ops import all_reduce +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.layers.activation import SiluAndMul +from tokenspeed.runtime.layers.layernorm import RMSNorm +from tokenspeed.runtime.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput +from tokenspeed.runtime.layers.paged_attention import PagedAttention +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.layers.rotary_embedding import get_rope +from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader +from tokenspeed.runtime.models.utils import validate_attention_partition +from tokenspeed.runtime.utils import add_prefix +from tokenspeed.runtime.utils.env import global_server_args_dict + + +class DFlashAttention(nn.Module): + def __init__( + self, + config, + mapping: Mapping, + layer_id: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.mapping = mapping + self.hidden_size = int(config.hidden_size) + self.tp_rank = self.mapping.attn.tp_rank + self.tp_size = self.mapping.attn.tp_size + self.total_num_heads = int(config.num_attention_heads) + self.total_num_kv_heads = int( + getattr(config, "num_key_value_heads", self.total_num_heads) + ) + validate_attention_partition( + self.total_num_heads, + self.total_num_kv_heads, + self.tp_size, + ) + self.num_heads = self.total_num_heads // self.tp_size + self.num_kv_heads = max(1, self.total_num_kv_heads // self.tp_size) + self.head_dim = int( + getattr(config, "head_dim", self.hidden_size // self.total_num_heads) + ) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + + self.qkv_proj = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=bool(getattr(config, "attention_bias", False)), + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + self.hidden_size, + bias=bool(getattr(config, "attention_bias", False)), + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + reduce_results=False, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + eps = float(getattr(config, "rms_norm_eps", 1e-6)) + self.q_norm = RMSNorm(self.head_dim, eps=eps) + self.k_norm = RMSNorm(self.head_dim, eps=eps) + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=int(getattr(config, "max_position_embeddings", 32768)), + base=float(getattr(config, "rope_theta", 1000000)), + rope_scaling=getattr(config, "rope_scaling", None), + ) + + # The FA4 MHA extend selector currently has no sliding-window kernel + # for this draft shape. Use full attention for draft proposals; target + # verification remains authoritative for accepted tokens. + sliding_window = -1 + self.attn = PagedAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + sliding_window_size=sliding_window, + ) + self.attn.non_causal = True + + def _apply_qk_norm( + self, q: torch.Tensor, k: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + q = q.reshape(-1, self.head_dim) + k = k.reshape(-1, self.head_dim) + q = self.q_norm(q).view(-1, self.q_size) + k = self.k_norm(k).view(-1, self.kv_size) + return q, k + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self._apply_qk_norm(q, k) + q, k = self.rotary_emb(positions, q, k) + k_cache = k.view(-1, self.num_kv_heads, self.head_dim) + v_cache = v.view(-1, self.num_kv_heads, self.head_dim) + ctx.token_to_kv_pool.set_kv_buffer( + self.attn, + out_cache_loc, + k_cache, + v_cache, + self.attn.k_scale, + self.attn.v_scale, + ) + attn_output = self.attn( + q, + None, + None, + ctx, + out_cache_loc, + save_kv_cache=False, + ) + if len(attn_output.size()) == 3: + attn_output = attn_output.reshape(attn_output.shape[0], -1) + output, _ = self.o_proj(attn_output) + return output + + def kv_proj_only( + self, hidden_states: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + qkv, _ = self.qkv_proj(hidden_states) + _, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + return k, v + + def apply_k_norm(self, k: torch.Tensor) -> torch.Tensor: + k_shape = k.shape + return self.k_norm(k.reshape(-1, self.head_dim)).view(k_shape) + + def apply_k_rope(self, positions: torch.Tensor, k: torch.Tensor) -> torch.Tensor: + dummy_q = k.new_empty(k.shape) + _, k = self.rotary_emb(positions, dummy_q, k) + return k + + +class DFlashMLP(nn.Module): + def __init__( + self, + config, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + hidden_size = int(config.hidden_size) + intermediate_size = int(config.intermediate_size) + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + tp_rank=mapping.dense.tp_rank, + tp_size=mapping.dense.tp_size, + tp_group=mapping.dense.tp_group, + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("down_proj", prefix), + reduce_results=False, + tp_rank=mapping.dense.tp_rank, + tp_size=mapping.dense.tp_size, + tp_group=mapping.dense.tp_group, + ) + if getattr(config, "hidden_act", "silu") != "silu": + raise ValueError("DFlash only supports silu activation.") + self.act_fn = SiluAndMul() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class DFlashDecoderLayer(nn.Module): + def __init__( + self, + config, + mapping: Mapping, + layer_id: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + hidden_size = int(config.hidden_size) + eps = float(getattr(config, "rms_norm_eps", 1e-6)) + self.mapping = mapping + self.input_layernorm = RMSNorm(hidden_size, eps=eps) + self.self_attn = DFlashAttention( + config=config, + mapping=mapping, + layer_id=layer_id, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + self.post_attention_layernorm = RMSNorm(hidden_size, eps=eps) + self.mlp = DFlashMLP( + config=config, + mapping=mapping, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if ctx.forward_mode.is_idle(): + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + elif ( + ctx.input_num_tokens > global_server_args_dict["comm_fusion_max_num_tokens"] + ): + hidden_states = all_reduce(hidden_states, self.mapping.dense.tp_group) + hidden_states, residual = self.input_layernorm(hidden_states, residual) + else: + hidden_states, residual, *_ = ( + self.input_layernorm.forward_with_allreduce_fusion( + self.mapping.dense.tp_rank, + self.mapping.dense.tp_group, + hidden_states, + residual, + ) + ) + + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ctx=ctx, + out_cache_loc=out_cache_loc, + ) + + if ctx.input_num_tokens > global_server_args_dict["comm_fusion_max_num_tokens"]: + hidden_states = all_reduce(hidden_states, self.mapping.attn.tp_group) + hidden_states, residual = self.post_attention_layernorm( + hidden_states, residual + ) + else: + hidden_states, residual, *_ = ( + self.post_attention_layernorm.forward_with_allreduce_fusion( + self.mapping.attn.tp_rank, + self.mapping.attn.tp_group, + hidden_states, + residual, + ) + ) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class DFlashDraftModel(nn.Module): + def __init__( + self, + config, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.mapping = mapping + eps = float(getattr(config, "rms_norm_eps", 1e-6)) + self.layers = nn.ModuleList( + [ + DFlashDecoderLayer( + config=config, + mapping=mapping, + layer_id=i, + quant_config=quant_config, + prefix=add_prefix(f"layers.{i}", prefix), + ) + for i in range(int(config.num_hidden_layers)) + ] + ) + self.norm = RMSNorm(int(config.hidden_size), eps=eps) + target_layer_ids = (getattr(config, "dflash_config", {}) or {}).get( + "target_layer_ids", [] + ) + self.num_context_features = len(target_layer_ids) + self.fc = nn.Linear( + self.num_context_features * int(config.hidden_size), + int(config.hidden_size), + bias=False, + ) + self.hidden_norm = RMSNorm(int(config.hidden_size), eps=eps) + self.block_size = int(getattr(config, "block_size", 8)) + + def project_target_hidden(self, target_hidden: torch.Tensor) -> torch.Tensor: + return self.hidden_norm(self.fc(target_hidden)) + + @torch.no_grad() + def forward( + self, + ctx: ForwardContext, + input_ids: torch.Tensor, + positions: torch.Tensor, + out_cache_loc: torch.Tensor, + input_lengths: torch.Tensor | None = None, + input_embeds: torch.Tensor | None = None, + **kwargs, + ) -> LogitsProcessorOutput: + if input_embeds is None: + if not ctx.forward_mode.is_idle(): + raise ValueError("DFlashDraftModel requires input_embeds.") + hidden_states = self.fc.weight.new_empty((0, int(self.config.hidden_size))) + else: + hidden_states = input_embeds + residual = None + + for layer in self.layers: + hidden_states, residual = layer( + positions=positions, + hidden_states=hidden_states, + ctx=ctx, + out_cache_loc=out_cache_loc, + residual=residual, + ) + + if residual is None: + hidden_states = self.norm(hidden_states) + else: + hidden_states, _ = self.norm(hidden_states, residual) + + return LogitsProcessorOutput( + next_token_logits=None, hidden_states=hidden_states + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + params_dict = dict(self.named_parameters()) + + def resolve_name(name: str) -> str | None: + if name in params_dict: + return name + if name.startswith("model.") and name[len("model.") :] in params_dict: + return name[len("model.") :] + return None + + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + for param_name, weight_name, shard_id in stacked_params_mapping: + if f".{weight_name}." not in name: + continue + resolved = resolve_name(name.replace(weight_name, param_name)) + if resolved is None: + continue + param = params_dict[resolved] + param.weight_loader(param, loaded_weight, shard_id) + break + else: + resolved = resolve_name(name) + if resolved is None: + continue + param = params_dict[resolved] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + +EntryClass = [DFlashDraftModel] diff --git a/python/tokenspeed/runtime/models/extensible.py b/python/tokenspeed/runtime/models/extensible.py new file mode 100755 index 0000000..743bc41 --- /dev/null +++ b/python/tokenspeed/runtime/models/extensible.py @@ -0,0 +1,196 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Extensible wrappers for injecting custom input and output processors.""" + +import importlib +from typing import Any + +from torch import Tensor, nn + +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.layers.logits_processor import ( + LogitsMetadata, + LogitsProcessorOutput, +) +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +# Used for Input/Output Processor sharing +class ContextBase: + """Base shared context for extensible input and output processors.""" + + def __init__(self, base_lm, config_dict: dict[str, Any]): + pass + + +class InputProcessorBase(nn.Module): + """Default input processor that falls back to token embeddings.""" + + def __init__(self, base_lm, ctx, config_dict: dict[str, Any]): + super().__init__() + self.base_lm = base_lm + + def forward( + self, + input_ids: Tensor, + positions: Tensor, + ctx: ForwardContext, + out_cache_loc: Tensor, + input_embeds: Tensor = None, + ) -> Tensor: + if input_embeds is not None: + return input_embeds + return self.base_lm.model.embed_tokens(input_ids) + + +class OutputProcessorBase(nn.Module): + """Default output processor that routes hidden states to logits.""" + + def __init__(self, base_lm, ctx, config_dict: dict[str, Any]): + super().__init__() + self.base_lm = base_lm + + def forward( + self, + input_ids: Tensor, + positions: Tensor, + ctx: ForwardContext, + output_hidden_states: Tensor, + ) -> LogitsProcessorOutput: + logits_metadata = LogitsMetadata.from_forward_context(ctx) + return self.base_lm.logits_processor( + input_ids, + output_hidden_states, + self.base_lm.lm_head, + logits_metadata, + ) + + +_EXT_CLS_REGISTRY: dict[str, type] = {} + + +def register_ext_cls(name: str, cls: type) -> None: + global _EXT_CLS_REGISTRY + _EXT_CLS_REGISTRY[name] = cls + + +def get_ext_cls(name: str) -> type: + if name not in _EXT_CLS_REGISTRY: + raise ValueError( + f"Input module {name} not found in registry. {_EXT_CLS_REGISTRY=}" + ) + return _EXT_CLS_REGISTRY[name] + + +register_ext_cls("ContextBase", ContextBase) +register_ext_cls("InputProcessorBase", InputProcessorBase) +register_ext_cls("OutputProcessorBase", OutputProcessorBase) + + +class ExtensibleLM(nn.Module): + """Wrap a base LM with pluggable context, input, and output processors.""" + + def __init__( + self, + base_lm: nn.Module, + ext_config: dict[str, Any], + ) -> None: + super().__init__() + self.base_lm = base_lm + + if "ext_def_file" in ext_config: + import os + import sys + from pathlib import Path + + ext_def_file = ext_config["ext_def_file"] + ext_def_dir = os.path.dirname(os.path.abspath(ext_def_file)) + sys.path.insert(0, ext_def_dir) + ext_def_module = f"{Path(ext_def_file).stem}" + logger.info( + "\x1b[32m[[ExtensibleLM] Loading ext_def_dir=%r, ext_def_module=%r]\x1b[0m", + ext_def_dir, + ext_def_module, + ) + importlib.import_module(ext_def_module) + + ctx_config = ext_config["context"] + ctx_name = ctx_config.pop("cls") + ctx_cls = get_ext_cls(ctx_name) + self.ctx: ContextBase = ctx_cls(base_lm, ctx_config) + + input_processor_config = ext_config["input_processor"] + input_processor_name = input_processor_config.pop("cls") + input_processor_cls = get_ext_cls(input_processor_name) + self.input_processor: InputProcessorBase = input_processor_cls( + self.base_lm, + self.ctx, + input_processor_config, + ).eval() + + output_processor_config = ext_config["output_processor"] + output_processor_name = output_processor_config.pop("cls") + output_processor_cls = get_ext_cls(output_processor_name) + self.output_processor: OutputProcessorBase = output_processor_cls( + self.base_lm, + self.ctx, + output_processor_config, + ).eval() + self.step = 0 + + @property + def logits_processor(self): + return self.base_lm.logits_processor + + @property + def lm_head(self): + return self.base_lm.lm_head + + def forward( + self, + ctx: ForwardContext, + input_ids: Tensor, + positions: Tensor, + out_cache_loc: Tensor, + input_embeds: Tensor = None, + ) -> LogitsProcessorOutput: + # input processor: get input hidden states + input_embeds = self.input_processor( + input_ids, positions, ctx, out_cache_loc, input_embeds + ) + + # base model forward + out_hidden_states, _ = self.base_lm.model( + input_ids=None, + positions=positions, + ctx=ctx, + out_cache_loc=out_cache_loc, + input_embeds=input_embeds, + ) + + # output processor: lm hidden states to logits + logits_output: LogitsProcessorOutput = self.output_processor( + input_ids, positions, ctx, out_hidden_states + ) + self.step += 1 + return logits_output diff --git a/python/tokenspeed/runtime/models/glm5.py b/python/tokenspeed/runtime/models/glm5.py new file mode 100644 index 0000000..b24ebe7 --- /dev/null +++ b/python/tokenspeed/runtime/models/glm5.py @@ -0,0 +1,1403 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Inference-only GLM 5 model.""" + +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import dataclass, replace +from typing import Any + +import torch +from tokenspeed_kernel.ops.attention import ( + dsa_decode_topk, + dsa_prefill_topk, +) +from torch import nn +from transformers import PretrainedConfig + +from tokenspeed.runtime.configs.utils import get_rope_theta +from tokenspeed.runtime.distributed import Mapping +from tokenspeed.runtime.distributed.comm_manager import CommManager +from tokenspeed.runtime.execution.breakable_cuda_graph import ( + break_point, + current_forward_ctx, + slice_to_real_tokens, +) +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.execution.forward_batch_info import ForwardMode +from tokenspeed.runtime.layers.layernorm import FusedRMSNorm, LayerNorm, RMSNorm +from tokenspeed.runtime.layers.linear import ( + MergedColumnParallelLinear, + ReplicatedLinear, +) +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.layers.quantization.utils import block_dequant +from tokenspeed.runtime.layers.rotary_embedding import get_rope +from tokenspeed.runtime.layers.vocab_parallel_embedding import VocabParallelEmbedding +from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader +from tokenspeed.runtime.models.deepseek_v3 import ( + DeepseekV3AttentionMLA, + DeepseekV3DecoderLayer, + DeepseekV3ForCausalLM, + DeepseekV3MLP, + DeepseekV3Model, + DeepseekV3MoE, + get_layer_id, +) +from tokenspeed.runtime.utils import add_prefix +from tokenspeed.runtime.utils.env import global_server_args_dict + +_INDEXER_PREFILL_MAX_LOGITS_MB_ARG = "deepseek_v4_indexer_prefill_max_logits_mb" + + +@dataclass +class GlmDsaIndexerOutput: + query: torch.Tensor + key: torch.Tensor + weights: torch.Tensor + + +@dataclass +class GlmDsaPrefillTopK: + workspace_indices: torch.Tensor + topk_lens: torch.Tensor + block_tables: torch.Tensor + seq_lens: torch.Tensor + max_seq_len: int + kv_workspace_slots: torch.Tensor + + +@dataclass +class GlmDsaDecodeTopK: + topk_indices: torch.Tensor + topk_lens: torch.Tensor + + +@dataclass(frozen=True) +class GlmDsaDecodeWindow: + start: int + end: int + num_tokens: int + num_reqs: int + q_len_per_req: int + + +def _glm_dsa_skip_indexer_topk(config, layer_id: int | None) -> bool: + if layer_id is None: + return False + indexer_types = getattr(config, "indexer_types", None) + if indexer_types is not None and layer_id < len(indexer_types): + return indexer_types[layer_id] in ("S", "shared") + pattern = getattr(config, "index_topk_pattern", None) + if pattern is not None and layer_id < len(pattern): + return pattern[layer_id] in ("S", "shared") + freq = int(getattr(config, "index_topk_freq", 1) or 1) + if freq <= 1: + return False + offset = getattr(config, "index_skip_topk_offset", None) + if offset is None: + return max(layer_id - 1, 0) % freq != 0 + if offset <= 0: + raise ValueError( + "index_skip_topk_offset must be positive; offset <= 0 marks " + "layer 0 as shared with no prior top-k to reuse" + ) + return max(layer_id - offset + 1, 0) % freq != 0 + + +def _build_prefill_kv_workspace_slots( + *, + block_tables: torch.Tensor, + seq_lens: torch.Tensor, + max_seq_len: int, + page_size: int, + device: torch.device, +) -> torch.Tensor: + local_offsets = torch.arange( + int(max_seq_len), + dtype=torch.int64, + device=device, + ) + page_offsets = torch.div( + local_offsets, + int(page_size), + rounding_mode="floor", + ) + block_offsets = local_offsets % int(page_size) + pages = block_tables.to(device=device, dtype=torch.int64).index_select( + 1, + page_offsets, + ) + slots = pages * int(page_size) + block_offsets + valid = local_offsets.unsqueeze(0) < seq_lens.to( + device=device, + dtype=torch.int64, + ).unsqueeze(1) + return slots[valid].contiguous() + + +def _glm_dsa_rope_scaling( + rope_scaling: dict[str, Any] | None, +) -> dict[str, Any] | None: + if not rope_scaling or "factor" not in rope_scaling: + return None + + rope_scaling = dict(rope_scaling) + rope_scaling["rope_type"] = "deepseek_yarn" + return rope_scaling + + +class GlmDsaIndexer(nn.Module): + def __init__( + self, + config: PretrainedConfig, + hidden_size: int, + q_lora_rank: int, + qk_rope_head_dim: int, + rope_theta: float, + rope_scaling: dict[str, Any] | None, + max_position_embeddings: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.index_topk = config.index_topk + self.index_n_heads = config.index_n_heads + self.index_head_dim = config.index_head_dim + self.rope_head_dim = int(qk_rope_head_dim) + self.softmax_scale = self.index_head_dim**-0.5 + + if self.rope_head_dim <= 0 or self.rope_head_dim > self.index_head_dim: + raise ValueError( + "GLM DSA indexer requires 0 < qk_rope_head_dim <= index_head_dim; " + f"got qk_rope_head_dim={self.rope_head_dim}, " + f"index_head_dim={self.index_head_dim}" + ) + + self.wq_b = ReplicatedLinear( + q_lora_rank, + self.index_n_heads * self.index_head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("wq_b", prefix), + ) + self.wk = ReplicatedLinear( + hidden_size, + self.index_head_dim, + bias=False, + quant_config=quant_config, + prefix=add_prefix("wk", prefix), + ) + self.weights_proj = ReplicatedLinear( + hidden_size, + self.index_n_heads, + bias=False, + quant_config=None, + prefix=add_prefix("weights_proj", prefix), + ) + self.wk_weights_proj = MergedColumnParallelLinear( + hidden_size, + [self.index_head_dim, self.index_n_heads], + bias=False, + quant_config=None, + prefix=add_prefix("wk_weights_proj", prefix), + ) + self._wk_weights_proj_loaded = False + self.k_norm = LayerNorm(self.index_head_dim, eps=1e-6) + + rope_scaling = _glm_dsa_rope_scaling(rope_scaling) + self.rotary_emb = get_rope( + self.rope_head_dim, + rotary_dim=self.rope_head_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + is_neox_style=not getattr(config, "indexer_rope_interleave", False), + ) + + def set_wk_weights_proj_loaded(self, loaded: bool = True) -> None: + self._wk_weights_proj_loaded = bool(loaded) + + def forward( + self, + hidden_states: torch.Tensor, + q_lora: torch.Tensor, + positions: torch.Tensor, + ) -> GlmDsaIndexerOutput: + index_q = self.wq_b(q_lora)[0] + index_q = index_q.view(-1, self.index_n_heads, self.index_head_dim) + if self._wk_weights_proj_loaded: + key_weights = self.wk_weights_proj(hidden_states)[0] + index_k, weights = key_weights.split( + [self.index_head_dim, self.index_n_heads], + dim=-1, + ) + else: + index_k = self.wk(hidden_states)[0] + weights = self.weights_proj(hidden_states)[0] + index_k = self.k_norm(index_k) + + q_rope, k_rope = self.rotary_emb( + positions, + index_q[..., : self.rope_head_dim], + index_k[:, None, : self.rope_head_dim], + ) + # Noops if the RoPE is in-place applied + index_q[..., : self.rope_head_dim] = q_rope + index_k[:, : self.rope_head_dim] = k_rope.squeeze(1) + + return GlmDsaIndexerOutput( + query=index_q, + key=index_k, + weights=weights.float() * (self.index_n_heads**-0.5), + ) + + +class GlmMoeDsaAttention(DeepseekV3AttentionMLA): + _MLA_KERNEL_BACKENDS = ("trtllm_mla", "tokenspeed_mla", "dsa") + _RAGGED_PREFILL_BACKENDS = ("trtllm_mla", "tokenspeed_mla", "dsa") + rope_is_neox_style = False + + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + hidden_size: int, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + q_lora_rank: int, + kv_lora_rank: int, + rope_theta: float = 10000, + rope_scaling: dict[str, Any] | None = None, + max_position_embeddings: int = 8192, + quant_config: QuantizationConfig | None = None, + layer_id=None, + prefix: str = "", + reduce_attn_results=True, + alt_stream: torch.cuda.Stream | None = None, + skip_rope: bool = False, + is_nextn: bool = False, + ) -> None: + rope_scaling = _glm_dsa_rope_scaling(rope_scaling) + super().__init__( + config=config, + mapping=mapping, + hidden_size=hidden_size, + num_heads=num_heads, + qk_nope_head_dim=qk_nope_head_dim, + qk_rope_head_dim=qk_rope_head_dim, + v_head_dim=v_head_dim, + q_lora_rank=q_lora_rank, + kv_lora_rank=kv_lora_rank, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + layer_id=layer_id, + prefix=prefix, + reduce_attn_results=reduce_attn_results, + alt_stream=alt_stream, + skip_rope=skip_rope, + ) + if q_lora_rank is None: + raise ValueError("GLM DSA requires q_lora_rank.") + # Let process_weights choose DeepGEMM only after it has transformed + # FP8 block scales into the layout that kernel expects. + self.q_a_layernorm = RMSNorm(q_lora_rank, eps=1e-6) + self.kv_a_layernorm = RMSNorm(kv_lora_rank, eps=1e-6) + self.fused_qk_layernorm = FusedRMSNorm( + self.q_a_layernorm, + self.kv_a_layernorm, + ) + self.index_topk = config.index_topk + self.is_nextn = is_nextn + # NextN/MTP has its own indexer weights but may reuse the previous + # draft iteration's top-k. Shared target layers do not have usable + # indexer weights and must consume the context-carried top-k. + self.skip_indexer_topk = ( + True if is_nextn else _glm_dsa_skip_indexer_topk(config, layer_id) + ) + if self.skip_indexer_topk and not self.is_nextn: + self.indexer = None + else: + self.indexer = GlmDsaIndexer( + config=config, + hidden_size=hidden_size, + q_lora_rank=q_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + prefix=add_prefix("indexer", prefix), + ) + self._decode_topk_indices_buffer: torch.Tensor | None = None + self._decode_topk_lens_buffer: torch.Tensor | None = None + + def _get_decode_topk_workspace( + self, + attr_name: str, + rows: int, + cols: int, + device: torch.device, + fill_value: int | None = -1, + ) -> torch.Tensor: + buffer = getattr(self, attr_name, None) + if ( + buffer is None + or buffer.device != device + or buffer.shape[0] < rows + or buffer.shape[1] != cols + ): + # A captured CUDA graph may still reference the old buffer; keep + # it alive so a regrow never frees memory a graph replays into. + if buffer is not None: + self._retire_decode_workspace(buffer) + buffer = torch.empty( + (rows, cols), + dtype=torch.int32, + device=device, + ) + setattr(self, attr_name, buffer) + workspace = buffer[:rows] + if fill_value is not None: + workspace.fill_(fill_value) + return workspace + + def _get_decode_topk_lens_workspace( + self, + rows: int, + device: torch.device, + ) -> torch.Tensor: + buffer = getattr(self, "_decode_topk_lens_buffer", None) + if buffer is None or buffer.device != device or buffer.numel() < rows: + if buffer is not None: + self._retire_decode_workspace(buffer) + buffer = torch.empty( + (rows,), + dtype=torch.int32, + device=device, + ) + self._decode_topk_lens_buffer = buffer + workspace = buffer[:rows] + workspace.fill_(0) + return workspace + + @staticmethod + def _resolve_decode_q_len( + ctx: ForwardContext, + num_decode_tokens: int, + num_decode_reqs: int, + ) -> int: + """Per-request query rows, derived from the actual batch shape. + + Spec-verify and the draft first step can both feed multiple query rows + per request, while the draft model's later decode steps feed one row. + The draft attention backend inherits the target verify width from the + shared config, so trust the actual input row count instead of backend + metadata. + """ + if num_decode_reqs > 0 and num_decode_tokens > 0: + q_len, rem = divmod(int(num_decode_tokens), int(num_decode_reqs)) + if rem == 0 and q_len > 0: + return q_len + return 1 + + @staticmethod + def _resolve_num_decode_tokens( + ctx: ForwardContext, + *, + total_tokens: int, + num_decode_reqs: int, + ) -> int: + if num_decode_reqs <= 0 or total_tokens <= 0: + return 0 + spec_width = int(getattr(ctx.attn_backend, "spec_num_tokens", 1) or 1) + expected_decode_tokens = num_decode_reqs * spec_width + return min(int(total_tokens), int(expected_decode_tokens)) + + @staticmethod + def _resolve_decode_req_count( + ctx: ForwardContext, + metadata: Any, + ) -> int: + num_extends = int(getattr(metadata, "num_extends", 0) or 0) + limits = [max(0, int(ctx.bs) - int(ctx.num_extends))] + + seq_lens = getattr(metadata, "seq_lens_k", None) + if seq_lens is not None: + limits.append(max(0, int(seq_lens.shape[0]) - num_extends)) + + block_tables = getattr(metadata, "block_kv_indices", None) + if block_tables is not None: + limits.append(max(0, int(block_tables.shape[0]) - num_extends)) + + return min(limits) + + @staticmethod + def _resolve_decode_window( + ctx: ForwardContext, + metadata: Any, + *, + total_tokens: int, + ) -> GlmDsaDecodeWindow: + num_decode_reqs = GlmMoeDsaAttention._resolve_decode_req_count(ctx, metadata) + num_decode_tokens = GlmMoeDsaAttention._resolve_num_decode_tokens( + ctx, + total_tokens=total_tokens, + num_decode_reqs=num_decode_reqs, + ) + if total_tokens < num_decode_tokens: + raise RuntimeError( + "GLM DSA decode token split is invalid: " + f"tokens={total_tokens}, decode_tokens={num_decode_tokens}" + ) + q_len_per_req = GlmMoeDsaAttention._resolve_decode_q_len( + ctx, num_decode_tokens, num_decode_reqs + ) + decode_start = int(total_tokens) - int(num_decode_tokens) + return GlmDsaDecodeWindow( + start=decode_start, + end=decode_start + int(num_decode_tokens), + num_tokens=int(num_decode_tokens), + num_reqs=int(num_decode_reqs), + q_len_per_req=int(q_len_per_req), + ) + + @staticmethod + def _slice_decode_topk( + decode_topk: GlmDsaDecodeTopK, + start: int, + end: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + return decode_topk.topk_indices[start:end], decode_topk.topk_lens[start:end] + + def _retire_decode_workspace(self, buffer: torch.Tensor) -> None: + retired = getattr(self, "_retired_decode_workspaces", None) + if retired is None: + retired = [] + self._retired_decode_workspaces = retired + retired.append(buffer) + + @staticmethod + def _check_decode_q_len_per_req(q_len_per_req: int) -> None: + # Multi-step MTP verify runs num_draft_tokens query rows per request. + # DeepGEMM paged MQA logits (our fork) and FlashMLA sparse decode are + # both verified bit-exact against batch expansion up to next_n = 6, + # which covers --speculative-num-steps 5 (5 draft + 1 bonus). + if not 1 <= q_len_per_req <= 6: + raise NotImplementedError( + "GLM DSA sparse decode supports 1-6 query tokens per request " + f"(verified next_n <= 6), got {q_len_per_req}." + ) + + def _compute_decode_topk_indices( + self, + indexer_output: GlmDsaIndexerOutput, + ctx: ForwardContext, + ) -> GlmDsaDecodeTopK | None: + metadata = getattr(ctx.attn_backend, "forward_decode_metadata", None) + if metadata is None or metadata.block_kv_indices is None: + return None + num_tokens = indexer_output.query.shape[0] + decode_window = self._resolve_decode_window( + ctx, metadata, total_tokens=num_tokens + ) + if decode_window.num_reqs <= 0 or num_tokens == 0: + return None + self._check_decode_q_len_per_req(decode_window.q_len_per_req) + + num_extends = int(metadata.num_extends or 0) + seq_lens = metadata.seq_lens_k[ + num_extends : num_extends + decode_window.num_reqs + ] + if seq_lens.numel() == 0: + return None + + block_tables = metadata.block_kv_indices[ + num_extends : num_extends + decode_window.num_reqs + ] + topk = self.index_topk + return self._compute_decode_topk_indices_portable( + indexer_output=indexer_output, + ctx=ctx, + seq_lens=seq_lens, + block_tables=block_tables, + q_len_per_req=decode_window.q_len_per_req, + decode_start=decode_window.start, + num_tokens=num_tokens, + num_decode_tokens=decode_window.num_tokens, + topk=topk, + ) + + def _compute_decode_topk_indices_portable( + self, + *, + indexer_output: GlmDsaIndexerOutput, + ctx: ForwardContext, + seq_lens: torch.Tensor, + block_tables: torch.Tensor, + q_len_per_req: int, + decode_start: int, + num_tokens: int, + num_decode_tokens: int, + topk: int, + ) -> GlmDsaDecodeTopK: + q = indexer_output.query[decode_start : decode_start + num_decode_tokens] + weights = indexer_output.weights[ + decode_start : decode_start + num_decode_tokens + ] + index_k_cache = ctx.token_to_kv_pool.get_index_k_buffer(self.attn_mqa.layer_id) + if index_k_cache is None: + raise RuntimeError("GLM DSA top-k requires an index-K cache.") + + topk_indices = self._get_decode_topk_workspace( + "_decode_topk_indices_buffer", + num_tokens, + topk, + q.device, + fill_value=-1, + ) + topk_slice = topk_indices[decode_start : decode_start + num_decode_tokens] + topk_lens = self._get_decode_topk_lens_workspace(num_tokens, q.device) + topk_lens_slice = topk_lens[decode_start : decode_start + num_decode_tokens] + + metadata = ctx.attn_backend.forward_decode_metadata + seq_lens_2d = ( + metadata._dsa_seq_lens_2d[ctx.num_extends :] + if q_len_per_req > 1 + else seq_lens.unsqueeze(1) + ) + dsa_decode_topk( + q, + weights, + seq_lens, + block_tables, + page_size=ctx.token_to_kv_pool.page_size, + topk=topk, + softmax_scale=self.indexer.softmax_scale, + q_len_per_req=q_len_per_req, + index_k_cache=index_k_cache, + seq_lens_2d=seq_lens_2d, + plan=metadata._dsa_plan, + out=topk_slice, + lens_out=topk_lens_slice, + ) + return GlmDsaDecodeTopK( + topk_indices=topk_indices, + topk_lens=topk_lens, + ) + + def _compute_prefill_topk_indices( + self, + indexer_output: GlmDsaIndexerOutput, + ctx: ForwardContext, + num_prefill_tokens: int, + ) -> GlmDsaPrefillTopK | None: + chunk_meta = ctx.attn_backend.chunked_prefill_metadata + prefix_lens = chunk_meta.extend_prefix_lens[: ctx.num_extends].to(torch.int32) + extend_lens = chunk_meta.extend_seq_lens[: ctx.num_extends].to(torch.int32) + seq_lens = prefix_lens + extend_lens + if seq_lens.numel() == 0: + return None + if int(extend_lens.sum().item()) != num_prefill_tokens: + raise RuntimeError( + "GLM DSA prefill token count mismatch: " + f"metadata={int(extend_lens.sum().item())}, " + f"tokens={num_prefill_tokens}" + ) + if ctx.req_to_page is None: + raise RuntimeError("GLM DSA sparse prefill requires req_to_page metadata") + + topk = self.index_topk + page_size = ctx.token_to_kv_pool.page_size + max_seq_len = int(seq_lens.max().item()) + max_pages = (max_seq_len + page_size - 1) // page_size + block_tables = chunk_meta.block_tables[:, :max_pages].to( + device=indexer_output.query.device, + dtype=torch.int32, + ) + kv_workspace_slots = _build_prefill_kv_workspace_slots( + block_tables=block_tables, + seq_lens=seq_lens, + max_seq_len=max_seq_len, + page_size=page_size, + device=indexer_output.query.device, + ) + return self._compute_prefill_topk_indices_portable( + indexer_output=indexer_output, + ctx=ctx, + prefix_lens=prefix_lens, + extend_lens=extend_lens, + seq_lens=seq_lens, + block_tables=block_tables, + kv_workspace_slots=kv_workspace_slots, + max_seq_len=max_seq_len, + num_prefill_tokens=num_prefill_tokens, + topk=topk, + ) + + def _compute_prefill_topk_indices_portable( + self, + *, + indexer_output: GlmDsaIndexerOutput, + ctx: ForwardContext, + prefix_lens: torch.Tensor, + extend_lens: torch.Tensor, + seq_lens: torch.Tensor, + block_tables: torch.Tensor, + kv_workspace_slots: torch.Tensor, + max_seq_len: int, + num_prefill_tokens: int, + topk: int, + ) -> GlmDsaPrefillTopK: + q = indexer_output.query[:num_prefill_tokens].contiguous() + weights = indexer_output.weights[:num_prefill_tokens].float().contiguous() + + req_ids = torch.arange( + seq_lens.numel(), + dtype=torch.int64, + device=q.device, + ) + token_req = torch.repeat_interleave(req_ids, extend_lens.to(torch.int64)) + extend_cu = torch.zeros( + extend_lens.numel() + 1, + dtype=torch.int64, + device=q.device, + ) + torch.cumsum(extend_lens.to(torch.int64), dim=0, out=extend_cu[1:]) + token_offsets = torch.arange( + num_prefill_tokens, dtype=torch.int64, device=q.device + ) - extend_cu.index_select(0, token_req) + causal_lens = ( + prefix_lens.to(torch.int64).index_select(0, token_req) + token_offsets + 1 + ) + seq_cu = torch.zeros( + seq_lens.numel() + 1, + dtype=torch.int64, + device=q.device, + ) + torch.cumsum(seq_lens.to(torch.int64), dim=0, out=seq_cu[1:]) + row_starts = seq_cu.index_select(0, token_req) + row_ends = row_starts + causal_lens + + index_k_cache = ctx.token_to_kv_pool.get_index_k_buffer(self.attn_mqa.layer_id) + if index_k_cache is None: + raise RuntimeError("GLM DSA top-k requires an index-K cache.") + + max_logits_mb = int(global_server_args_dict[_INDEXER_PREFILL_MAX_LOGITS_MB_ARG]) + workspace_indices, topk_lens = dsa_prefill_topk( + q, + weights, + kv_workspace_slots, + row_starts.to(torch.int32).contiguous(), + row_ends.to(torch.int32).contiguous(), + topk=topk, + softmax_scale=self.indexer.softmax_scale, + index_k_cache=index_k_cache, + page_size=ctx.token_to_kv_pool.page_size, + max_logits_bytes=max(1, max_logits_mb) * 1024 * 1024, + ) + return GlmDsaPrefillTopK( + workspace_indices=workspace_indices, + topk_lens=topk_lens, + block_tables=block_tables, + seq_lens=seq_lens.to(device=q.device, dtype=torch.int32), + max_seq_len=max_seq_len, + kv_workspace_slots=kv_workspace_slots, + ) + + @break_point + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + comm_manager: CommManager, + block_scale: torch.Tensor | None = None, + ) -> torch.Tensor: + """GLM-5 DSA attention, one COARSE breakable-graph break point. + + Like DeepSeek-V4 it does paged-cache writes, a data-dependent indexer + -> top-k stage and the FlashMLA sparse kernel (plus pre-attn + collectives), none capturable. Under a prefill-graph capture the whole + attention runs eager (reading the live ``ctx``) while the layer's + norms + MoE stay graphed; direct call otherwise (see ``break_point``). + Padded token-shaped inputs are sliced to the real count the live + metadata describes -- DSA and the decode-window split (which derives + ``decode_start`` from the total token count) must not see padded rows, + or decode rows get sliced out of the padded tail. Mirrors the + DeepSeek-V4 DSA break. + """ + # Empty (idle / DP-idle) batch: explicit skip, like the sibling MLP/MoE forwards. + if hidden_states.shape[0] == 0: + return hidden_states + qkv = self.fused_qkv_a_proj_with_mqa( + hidden_states, + block_scale, + torch.bfloat16, + ) + # The fused QKV-A weight may be zero-padded on its output dim to a + # multiple of 128 so the FP8 block-scale GEMM stays numerically valid + # (see GlmMoeDsaForCausalLM._pad_fused_qkv_a_proj_for_fp8_blockscale). + # Drop the padding columns before the split / comm. No-op when the + # projection output already matches the logical width. + _qkv_width = self.q_lora_rank + self.kv_lora_rank + self.qk_rope_head_dim + if qkv.shape[-1] != _qkv_width: + qkv = qkv[..., :_qkv_width] + qkv = comm_manager.pre_attn_comm(qkv, ctx) + # Slice only under a breakable capture/replay (see the DeepSeek-V4 break): + # eager forwards (incl. MTP draft steps) are never padded. Sliced AFTER + # the pre-attn comm: at replay ``_padded_to`` pins ``global_num_tokens`` + # to the padded bucket, so the comm must see padded-length rows; only + # the DSA stack below needs exactly the real rows. + _metadata = getattr(ctx.attn_backend, "forward_metadata", None) + _token_to_req = getattr(_metadata, "token_to_req_indices", None) + if current_forward_ctx() is not None and _token_to_req is not None: + positions, qkv, out_cache_loc = slice_to_real_tokens( + _token_to_req.numel(), positions, qkv, out_cache_loc + ) + q_a, latent_cache = qkv.split( + [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], + dim=-1, + ) + kv_a = latent_cache[..., : self.kv_lora_rank] + q_norm = torch.empty_like(q_a) + if q_a.size(0) > 0: + self.fused_qk_layernorm(input_q_a=q_a, input_kv_a=kv_a, output_q_a=q_norm) + + decode_metadata = getattr(ctx.attn_backend, "forward_decode_metadata", None) + num_attn_tokens = int(q_norm.shape[0]) + decode_window = self._resolve_decode_window( + ctx, + decode_metadata, + total_tokens=num_attn_tokens, + ) + num_decode_tokens = decode_window.num_tokens + num_prefill_tokens = decode_window.start + decode_start = decode_window.start + decode_end = decode_window.end + + should_compute_indexer = not self.skip_indexer_topk or ( + self.is_nextn + and ( + (num_prefill_tokens > 0 and ctx.dsa_prefill_topk is None) + or (num_decode_tokens > 0 and ctx.dsa_decode_topk is None) + ) + ) + if should_compute_indexer: + hidden_states = comm_manager.pre_attn_comm(hidden_states, ctx) + indexer_output = self.indexer(hidden_states, q_norm, positions) + ctx.token_to_kv_pool.set_index_k_buffer( + self.attn_mqa.layer_id, + out_cache_loc, + indexer_output.key, + ) + if ctx.num_extends > 0: + ctx.dsa_prefill_topk = self._compute_prefill_topk_indices( + indexer_output, + ctx, + num_prefill_tokens, + ) + if ctx.num_extends < ctx.bs: + ctx.dsa_decode_topk = self._compute_decode_topk_indices( + indexer_output, + ctx, + ) + + q = self.q_b_proj(q_norm)[0] + attn_output = torch.empty( + q.size(0), + self.num_local_heads * self.v_head_dim, + dtype=q.dtype, + device=q.device, + ) + + if ctx.num_extends > 0: + prefill_ctx = replace( + ctx, + bs=ctx.num_extends, + input_num_tokens=num_prefill_tokens, + forward_mode=ForwardMode.EXTEND, + ) + if ctx.dsa_prefill_topk is None: + raise RuntimeError( + "GLM DSA sparse prefill requires computed top-k indices." + ) + self.forward_dsa_sparse_prefill( + positions[:num_prefill_tokens], + q[:num_prefill_tokens], + latent_cache[:num_prefill_tokens], + prefill_ctx, + out_cache_loc[:num_prefill_tokens], + attn_output[:num_prefill_tokens], + prefill_topk=ctx.dsa_prefill_topk, + ) + + if num_decode_tokens > 0: + decode_ctx = replace( + ctx, + bs=decode_window.num_reqs, + num_extends=0, + input_num_tokens=num_decode_tokens, + forward_mode=ForwardMode.DECODE, + ) + if ctx.dsa_decode_topk is None: + raise RuntimeError( + "GLM DSA sparse decode requires computed top-k indices." + ) + topk_indices, topk_lens = self._slice_decode_topk( + ctx.dsa_decode_topk, + decode_start, + decode_end, + ) + self.forward_absorb( + positions[decode_start:decode_end], + q[decode_start:decode_end], + latent_cache[decode_start:decode_end], + decode_ctx, + out_cache_loc[decode_start:decode_end], + attn_output[decode_start:decode_end], + topk_indices=topk_indices, + topk_lens=topk_lens, + ) + + if ctx.accept_lengths is not None: + attn_output = attn_output.index_select(0, ctx.gather_ids) + output, _ = self.o_proj(attn_output) + return output + + def forward_dsa_sparse_prefill( + self, + positions: torch.Tensor, + q: torch.Tensor, + latent_cache: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + output: torch.Tensor, + *, + prefill_topk: GlmDsaPrefillTopK, + ) -> torch.Tensor: + Q, _ = self.forward_absorb_qkv_proj( + q, + latent_cache, + positions, + ctx, + out_cache_loc, + ) + attn_output = ctx.attn_backend.forward_sparse_prefill( + q=Q, + layer=self.attn_mqa, + token_to_kv_pool=ctx.token_to_kv_pool, + block_tables=prefill_topk.block_tables, + seq_lens=prefill_topk.seq_lens, + workspace_indices=prefill_topk.workspace_indices, + topk_lens=prefill_topk.topk_lens, + kv_workspace_slots=prefill_topk.kv_workspace_slots, + max_seq_len=prefill_topk.max_seq_len, + ) + attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank) + output_view = output.view(-1, self.num_local_heads, self.v_head_dim) + torch.bmm( + attn_output.transpose(0, 1), + self.w_vc, + out=output_view.transpose(0, 1), + ) + return output + + def forward_absorb( + self, + positions: torch.Tensor, + q: torch.Tensor, + latent_cache: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + output: torch.Tensor, + topk_indices: torch.Tensor | None = None, + topk_lens: torch.Tensor | None = None, + ) -> torch.Tensor: + Q, K = self.forward_absorb_qkv_proj( + q, + latent_cache, + positions, + ctx, + out_cache_loc, + ) + return self.forward_absorb_attn_v_proj( + Q, + K, + ctx, + out_cache_loc, + output, + topk_indices=topk_indices, + topk_lens=topk_lens, + ) + + def forward_absorb_attn_v_proj( + self, + Q, + K, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + output: torch.Tensor, + topk_indices: torch.Tensor | None = None, + topk_lens: torch.Tensor | None = None, + ) -> torch.Tensor: + need_save_kv = False + if self.attention_backend not in self._MLA_KERNEL_BACKENDS: + need_save_kv = not self.use_fused_set_kv_buffer + + attn_output = self.attn_mqa( + Q, + K, + K[..., : self.kv_lora_rank], + ctx, + out_cache_loc, + save_kv_cache=need_save_kv, + topk_indices=topk_indices, + topk_lens=topk_lens, + ) + attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank) + output_view = output.view(-1, self.num_local_heads, self.v_head_dim) + torch.bmm( + attn_output.transpose(0, 1), + self.w_vc, + out=output_view.transpose(0, 1), + ) + return output + + +class GlmMoeDsaDecoderLayer(DeepseekV3DecoderLayer): + def __init__( + self, + config: PretrainedConfig, + layer_id: int, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + is_nextn: bool = False, + prefix: str = "", + alt_stream: torch.cuda.Stream | None = None, + ) -> None: + nn.Module.__init__(self) + self.mapping = mapping + self.hidden_size = config.hidden_size + rope_theta = get_rope_theta(config) + rope_scaling = getattr(config, "rope_scaling", None) + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + + self.self_attn = GlmMoeDsaAttention( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + qk_nope_head_dim=config.qk_nope_head_dim, + qk_rope_head_dim=config.qk_rope_head_dim, + v_head_dim=config.v_head_dim, + q_lora_rank=( + config.q_lora_rank if hasattr(config, "q_lora_rank") else None + ), + kv_lora_rank=config.kv_lora_rank, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=( + None + if "self_attn" in getattr(config, "disable_quant_module", []) + else quant_config + ), + layer_id=layer_id, + prefix=add_prefix("self_attn", prefix), + reduce_attn_results=False, + alt_stream=alt_stream, + mapping=self.mapping, + is_nextn=is_nextn, + ) + + self.layer_id = layer_id + self.is_moe_layer = self._is_moe_layer(layer_id, is_nextn, config) + if self.is_moe_layer: + self.mlp = DeepseekV3MoE( + config=config, + mapping=self.mapping, + quant_config=quant_config, + layer_index=layer_id, + prefix=add_prefix("mlp", prefix), + alt_stream=alt_stream, + ) + else: + self.mlp = DeepseekV3MLP( + hidden_size=config.hidden_size, + intermediate_size=( + config.ffn_hidden_size + if hasattr(config, "ffn_hidden_size") + else config.intermediate_size + ), + hidden_act=config.hidden_act, + mapping=self.mapping, + quant_config=( + None + if "dense_mlp" in getattr(config, "disable_quant_module", []) + else quant_config + ), + prefix=add_prefix("mlp", prefix), + is_shared_expert=False, + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + self.comm_manager = CommManager( + mapping=self.mapping, + layer_id=self.layer_id, + is_moe=self.is_moe_layer, + prev_is_moe=self._is_moe_layer(layer_id - 1, is_nextn, config), + input_layernorm=self.input_layernorm, + post_attn_layernorm=self.post_attention_layernorm, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + residual: torch.Tensor | None, + ) -> torch.Tensor: + num_global_tokens, max_num_tokens_per_gpu = self.comm_manager.get_num_tokens( + ctx + ) + + if not ctx.forward_mode.is_idle(): + hidden_states, residual = self.comm_manager.input_reduce_norm( + hidden_states, residual + ) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ctx=ctx, + out_cache_loc=out_cache_loc, + comm_manager=self.comm_manager, + ) + if ctx.accept_lengths is not None: + residual = residual.index_select(0, ctx.gather_ids) + hidden_states, residual = self.comm_manager.post_attn_reduce_norm( + hidden_states, residual, ctx + ) + hidden_states = self.forward_mlp( + hidden_states, + residual, + ctx, + num_global_tokens, + max_num_tokens_per_gpu, + ) + else: + hidden_states = self.forward_mlp( + hidden_states, + residual, + ctx, + num_global_tokens, + max_num_tokens_per_gpu, + ) + return hidden_states, residual + + def forward_mlp( + self, + hidden_states, + residual, + ctx: ForwardContext, + num_global_tokens, + max_num_tokens_per_gpu, + ): + hidden_states = self.comm_manager.pre_mlp_comm(hidden_states, ctx) + if self.is_moe_layer: + hidden_states = self.mlp( + hidden_states, num_global_tokens, max_num_tokens_per_gpu + ) + else: + hidden_states = self.mlp(hidden_states) + hidden_states, residual = self.comm_manager.post_mlp_fused( + hidden_states, residual, ctx + ) + return hidden_states + + +class GlmMoeDsaModel(DeepseekV3Model): + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + nn.Module.__init__(self) + self.mapping = mapping + self.padding_id = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + ) + self.alt_stream = torch.cuda.Stream() + self.layers = nn.ModuleList( + [ + GlmMoeDsaDecoderLayer( + config, + layer_id, + mapping=self.mapping, + quant_config=quant_config, + prefix=add_prefix(f"layers.{layer_id}", prefix), + alt_stream=self.alt_stream, + ) + for layer_id in range(config.num_hidden_layers) + ] + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.layers_to_capture: set = set() + + +def pad_fused_qkv_a_proj_weight_for_fp8_blockscale(attn) -> None: + """Pad one attention module's fused QKV-A projection output dim to 128. + + The FP8 block-scale dense GEMM (deep_gemm / default ``mm`` path) returns NaN + when the output dim ``N`` is not a multiple of the 128 scale block. GLM-5.1's + fused QKV-A projection has ``N = q_lora_rank + kv_lora_rank + qk_rope_head_dim`` + (e.g. 2624), which is not 128-aligned, so attention output goes NaN. We + zero-pad the FP8 weight rows up to the next 128 multiple; + ``weight_scale_inv`` already has ``ceil(N/128)`` row blocks (covering the + padded rows) and the downstream ``qkv.split(...)`` drops the padding rows, so + real outputs are unchanged. No-op for bf16 weights or already-aligned ``N``. + + Shared by the main model (per decoder layer) and the NextN draft model (its + single DSA decoder), both of which carry the same fused QKV-A projection. + + Args: + attn: A GLM DSA attention module exposing ``fused_qkv_a_proj_with_mqa``. + """ + fp8_dtypes = (torch.float8_e4m3fn, getattr(torch, "float8_e4m3fnuz", None)) + fp8_dtypes = tuple(d for d in fp8_dtypes if d is not None) + proj = getattr(attn, "fused_qkv_a_proj_with_mqa", None) + weight = getattr(proj, "weight", None) + if weight is None or weight.dtype not in fp8_dtypes: + return + n = weight.shape[0] + if n % 128 == 0: + return + n_pad = ((n + 127) // 128) * 128 + pad = weight.new_zeros(n_pad - n, weight.shape[1]) + proj.weight = torch.nn.Parameter( + torch.cat([weight.data, pad], dim=0), requires_grad=False + ) + + +class GlmMoeDsaForCausalLM(DeepseekV3ForCausalLM): + model_cls = GlmMoeDsaModel + + def _record_fused_indexer_projection_shard( + self, + *, + module_name: str, + shard_id: int, + loaded_shards: dict[str, set[int]], + modules_dict: dict[str, nn.Module], + ) -> None: + shards = loaded_shards.setdefault(module_name, set()) + shards.add(int(shard_id)) + if shards != {0, 1}: + return + + module = modules_dict.get(module_name) + if isinstance(module, GlmDsaIndexer): + module.set_wk_weights_proj_loaded() + + def _load_fused_indexer_projection_shard( + self, + *, + module_name: str, + shard_id: int, + loaded_weight: torch.Tensor, + params_dict: dict[str, torch.Tensor], + modules_dict: dict[str, nn.Module], + loaded_shards: dict[str, set[int]], + ) -> bool: + param = params_dict.get(f"{module_name}.wk_weights_proj.weight") + if param is None: + return False + + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight, shard_id) + self._record_fused_indexer_projection_shard( + module_name=module_name, + shard_id=shard_id, + loaded_shards=loaded_shards, + modules_dict=modules_dict, + ) + return True + + def _flush_fused_indexer_fp8_wk( + self, + *, + module_name: str, + pending_fp8_wk: dict[str, dict[str, torch.Tensor]], + params_dict: dict[str, torch.Tensor], + modules_dict: dict[str, nn.Module], + loaded_shards: dict[str, set[int]], + ) -> None: + entry = pending_fp8_wk.get(module_name) + if not entry or "weight" not in entry or "scale" not in entry: + return + weight_block_size = getattr(self.quant_config, "weight_block_size", None) + if weight_block_size is None: + return + + weight_fp8 = entry["weight"] + scale = entry["scale"] + weight_bf16 = block_dequant( + weight_fp8, + scale, + list(weight_block_size), + ).to(torch.bfloat16) + if self._load_fused_indexer_projection_shard( + module_name=module_name, + shard_id=0, + loaded_weight=weight_bf16, + params_dict=params_dict, + modules_dict=modules_dict, + loaded_shards=loaded_shards, + ): + del pending_fp8_wk[module_name] + + def _try_load_fused_indexer_projection( + self, + *, + name: str, + loaded_weight: torch.Tensor, + params_dict: dict[str, torch.Tensor], + modules_dict: dict[str, nn.Module], + pending_fp8_wk: dict[str, dict[str, torch.Tensor]], + loaded_shards: dict[str, set[int]], + ) -> None: + if ".indexer.wk_weights_proj." in name: + return + + if ".indexer.weights_proj.weight" in name: + module_name = name.rsplit(".weights_proj.weight", 1)[0] + self._load_fused_indexer_projection_shard( + module_name=module_name, + shard_id=1, + loaded_weight=loaded_weight, + params_dict=params_dict, + modules_dict=modules_dict, + loaded_shards=loaded_shards, + ) + return + + if ".indexer.wk." not in name: + return + + module_name = name.rsplit(".wk.", 1)[0] + if name.endswith(".weight") and loaded_weight.dtype in ( + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + ): + pending_fp8_wk.setdefault(module_name, {})["weight"] = loaded_weight + self._flush_fused_indexer_fp8_wk( + module_name=module_name, + pending_fp8_wk=pending_fp8_wk, + params_dict=params_dict, + modules_dict=modules_dict, + loaded_shards=loaded_shards, + ) + return + + if name.endswith(".weight"): + self._load_fused_indexer_projection_shard( + module_name=module_name, + shard_id=0, + loaded_weight=loaded_weight, + params_dict=params_dict, + modules_dict=modules_dict, + loaded_shards=loaded_shards, + ) + return + + if "weight_scale_inv" in name: + pending_fp8_wk.setdefault(module_name, {})["scale"] = loaded_weight + self._flush_fused_indexer_fp8_wk( + module_name=module_name, + pending_fp8_wk=pending_fp8_wk, + params_dict=params_dict, + modules_dict=modules_dict, + loaded_shards=loaded_shards, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> None: + params_dict = dict(self.named_parameters()) + modules_dict = dict(self.named_modules()) + pending_fp8_wk: dict[str, dict[str, torch.Tensor]] = {} + loaded_fused_indexer_shards: dict[str, set[int]] = {} + + def base_weights(): + for name, loaded_weight in weights: + layer_id = get_layer_id(name) + if layer_id is not None and layer_id >= self.config.num_hidden_layers: + continue + if "rotary_emb.inv_freq" in name: + continue + if ".indexer." not in name: + yield name, loaded_weight + continue + + if name.endswith(".bias") and name not in params_dict: + continue + param = self.get_param(params_dict, name) + if param is None: + continue + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + self._try_load_fused_indexer_projection( + name=name, + loaded_weight=loaded_weight, + params_dict=params_dict, + modules_dict=modules_dict, + pending_fp8_wk=pending_fp8_wk, + loaded_shards=loaded_fused_indexer_shards, + ) + + super().load_weights(base_weights()) + self._pad_fused_qkv_a_proj_for_fp8_blockscale() + + def _pad_fused_qkv_a_proj_for_fp8_blockscale(self) -> None: + """Pad each decoder layer's fused QKV-A projection to a 128-multiple. + + See :func:`pad_fused_qkv_a_proj_weight_for_fp8_blockscale` for why this + is needed (FP8 block-scale GEMM returns NaN for non-128-aligned ``N``). + """ + for layer in getattr(self.model, "layers", []): + attn = getattr(layer, "self_attn", None) + if attn is not None: + pad_fused_qkv_a_proj_weight_for_fp8_blockscale(attn) + + +EntryClass = [GlmMoeDsaForCausalLM] diff --git a/python/tokenspeed/runtime/models/glm5_nextn.py b/python/tokenspeed/runtime/models/glm5_nextn.py new file mode 100644 index 0000000..6ede495 --- /dev/null +++ b/python/tokenspeed/runtime/models/glm5_nextn.py @@ -0,0 +1,508 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Inference-only GLM5 NextN speculative decoding.""" + +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import replace +from typing import Any + +import torch +from torch import nn +from transformers import PretrainedConfig + +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.execution.context import ( + ForwardContext, + report_collective_sizing, +) +from tokenspeed.runtime.layers.attention.dsa.utils import workspace_indices_to_kv_slots +from tokenspeed.runtime.layers.layernorm import RMSNorm +from tokenspeed.runtime.layers.linear import ReplicatedLinear +from tokenspeed.runtime.layers.logits_processor import LogitsMetadata, LogitsProcessor +from tokenspeed.runtime.layers.moe import ( + ExpertCheckpointSchema, + build_moe_checkpoint_loader, +) +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.layers.quantization.utils import block_dequant +from tokenspeed.runtime.layers.utils import ( + CP_METADATA, + ENABLE_CP, + cp_all_gather_rerange_output, + cp_split_and_rebuild_data, +) +from tokenspeed.runtime.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader +from tokenspeed.runtime.models.glm5 import ( + GlmMoeDsaDecoderLayer, + GlmMoeDsaForCausalLM, + pad_fused_qkv_a_proj_weight_for_fp8_blockscale, +) + +_NEXTN_SPEC_WEIGHT_NAMES = ( + "shared_head.norm", + "eh_proj", + "enorm", + "hnorm", +) + +_STACKED_PARAMS_MAPPING = ( + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), +) + + +class GlmMoeDsaModelNextN(nn.Module): + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + ) -> None: + super().__init__() + self.mapping = mapping + self.vocab_size = config.vocab_size + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + + self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.eh_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False) + + self.alt_stream = torch.cuda.Stream() + self.decoder = GlmMoeDsaDecoderLayer( + config, + 0, + mapping=self.mapping, + quant_config=quant_config, + is_nextn=True, + alt_stream=self.alt_stream, + ) + + self.shared_head = nn.Module() + self.shared_head.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + input_embeds: torch.Tensor | None = None, + captured_hidden_states: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + hidden_states = torch.where(positions.unsqueeze(-1) == 0, 0, hidden_states) + if captured_hidden_states is None: + if not ctx.forward_mode.is_idle(): + raise ValueError("GLM5 NextN requires captured_hidden_states.") + captured_hidden_states = hidden_states + + hidden_states = self.eh_proj( + torch.cat( + ( + self.enorm(hidden_states), + self.hnorm(captured_hidden_states), + ), + dim=-1, + ) + ) + + residual = None + if CP_METADATA: + hidden_states = cp_split_and_rebuild_data( + hidden_states, + CP_METADATA.value.split_list, + CP_METADATA.value.zigzag_index, + ) + positions = cp_split_and_rebuild_data( + positions, + CP_METADATA.value.split_list, + CP_METADATA.value.zigzag_index, + ) + hidden_states, residual = self.decoder( + positions, + hidden_states, + ctx, + out_cache_loc, + residual, + ) + + if not ctx.forward_mode.is_idle(): + if not ENABLE_CP: + hidden_states, _ = self.decoder.comm_manager.final_norm( + hidden_states, residual, ctx, self.shared_head.norm + ) + else: + hidden_states, _ = self.shared_head.norm(hidden_states, residual) + if CP_METADATA: + hidden_states = cp_all_gather_rerange_output( + hidden_states, + CP_METADATA.value, + self.mapping.attn.tp_rank, + self.mapping.attn.tp_group, + ) + return hidden_states, None + + +class GlmMoeDsaForCausalLMNextN(GlmMoeDsaForCausalLM): + compute_dsa_topk_first_step = True + + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + ) -> None: + nn.Module.__init__(self) + self.config = config + self.mapping = mapping + + if quant_config is not None and quant_config.get_name() == "nvfp4": + quant_config = None + + self.quant_config = quant_config + + self.model = GlmMoeDsaModelNextN( + config, mapping=self.mapping, quant_config=quant_config + ) + + if self.mapping.attn.has_dp: + self.lm_head = ReplicatedLinear( + config.hidden_size, + config.vocab_size, + bias=False, + ) + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + self.logits_processor = LogitsProcessor( + config, + skip_all_gather=self.mapping.attn.has_dp, + do_argmax=True, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + + @staticmethod + def _apply_first_step_correction(ctx: ForwardContext) -> None: + seq_lens_buf = ctx.draft_seq_lens_buf + accept_lengths = ctx.accept_lengths + if seq_lens_buf is None or accept_lengths is None: + return + num_extends = ctx.num_extends + if num_extends >= ctx.bs: + return + correction = ( + ctx.attn_backend.spec_num_tokens - accept_lengths[num_extends:] + ).to(seq_lens_buf.dtype) + seq_lens_buf[num_extends : ctx.bs].sub_(correction).clamp_(min=1) + ctx.attn_backend.advance_draft_forward_metadata(seq_lens_buf[: ctx.bs]) + + @staticmethod + def prepare_dsa_topk_for_mtp_decode( + dsa_topk: tuple[Any | None, Any | None], + gather_ids: torch.Tensor, + *, + num_prefill_rows: int = 0, + ) -> tuple[Any | None, Any | None]: + prefill_topk, decode_topk = dsa_topk + if decode_topk is None: + return dsa_topk + topk_indices = decode_topk.topk_indices + topk_lens = decode_topk.topk_lens + if topk_indices.shape[0] == 0: + return dsa_topk + if num_prefill_rows <= 0 and topk_indices.shape[0] <= gather_ids.numel(): + return dsa_topk + if num_prefill_rows <= 0: + selected_indices = topk_indices.index_select(0, gather_ids) + selected_lens = topk_lens.index_select(0, gather_ids) + else: + if prefill_topk is None: + return dsa_topk + num_prefill_rows = min(int(num_prefill_rows), gather_ids.numel()) + prefill_row_ids = gather_ids[:num_prefill_rows] + decode_row_ids = gather_ids[num_prefill_rows:] + selected_prefill_indices = workspace_indices_to_kv_slots( + prefill_topk.workspace_indices.index_select(0, prefill_row_ids), + prefill_topk.kv_workspace_slots, + ).to(device=topk_indices.device, dtype=topk_indices.dtype) + selected_prefill_lens = prefill_topk.topk_lens.index_select( + 0, + prefill_row_ids, + ).to( + device=topk_lens.device, + dtype=topk_lens.dtype, + ) + if decode_row_ids.numel() > 0: + selected_decode_indices = topk_indices.index_select(0, decode_row_ids) + selected_decode_lens = topk_lens.index_select(0, decode_row_ids) + selected_indices = torch.cat( + [selected_prefill_indices, selected_decode_indices], + dim=0, + ) + selected_lens = torch.cat( + [selected_prefill_lens, selected_decode_lens], + dim=0, + ) + else: + selected_indices = selected_prefill_indices + selected_lens = selected_prefill_lens + selected_decode_topk = replace( + decode_topk, + topk_indices=selected_indices, + topk_lens=selected_lens, + ) + return prefill_topk, selected_decode_topk + + @torch.no_grad() + def forward( + self, + ctx: ForwardContext, + input_ids: torch.Tensor, + positions: torch.Tensor, + out_cache_loc: torch.Tensor, + captured_hidden_states: torch.Tensor | None = None, + ) -> torch.Tensor: + with report_collective_sizing(ctx, ctx.bs, ctx.global_bs): + hidden_states, _ = self.model( + input_ids, + positions, + ctx, + out_cache_loc, + captured_hidden_states=captured_hidden_states, + ) + self._apply_first_step_correction(ctx) + logits_metadata = LogitsMetadata.from_forward_context(ctx) + return self.logits_processor( + input_ids, hidden_states, self.lm_head, logits_metadata + ) + + def get_hot_token_id(self) -> None: + return None + + def _nextn_layer_prefix(self, name: str) -> str | None: + if not hasattr(self.config, "num_nextn_predict_layers"): + raise ValueError("num_nextn_predict_layers is not in the config") + if self.config.num_nextn_predict_layers != 1: + raise ValueError("Only 1 nextn layer is supported") + + if self.config.num_nextn_predict_layers == self.config.num_hidden_layers: + prefix = "model.layers.0" + return prefix if name.startswith(prefix) else None + + if not name.startswith("model.layers."): + return None + name_parts = name.split(".") + if len(name_parts) < 3: + return None + try: + layer_id = int(name_parts[2]) + except ValueError: + return None + if layer_id < self.config.num_hidden_layers: + return None + return f"model.layers.{layer_id}" + + def _map_checkpoint_name(self, raw_name: str) -> str | None: + nextn_layer_prefix = self._nextn_layer_prefix(raw_name) + if nextn_layer_prefix is None: + return None + if "shared_head.head" in raw_name or "embed_tokens" in raw_name: + return None + if "rotary_emb.inv_freq" in raw_name: + return None + + if any(weight_name in raw_name for weight_name in _NEXTN_SPEC_WEIGHT_NAMES): + return raw_name.replace(nextn_layer_prefix, "model") + return raw_name.replace(nextn_layer_prefix, "model.decoder") + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> None: + fuse_qkv_a_proj = hasattr(self.config, "q_lora_rank") and ( + self.config.q_lora_rank is not None + ) + cached_a_proj: dict[str, torch.Tensor] | None = {} if fuse_qkv_a_proj else None + + params_dict = dict(self.named_parameters()) + modules_dict = dict(self.named_modules()) + pending_fp8_wk: dict[str, dict[str, torch.Tensor]] = {} + loaded_fused_indexer_shards: dict[str, set[int]] = {} + + moe_loader = build_moe_checkpoint_loader( + params_dict=params_dict, + expert_schema=ExpertCheckpointSchema( + gate_proj_name="gate_proj", + down_proj_name="down_proj", + up_proj_name="up_proj", + ), + num_experts=self.config.n_routed_experts, + ep_rank=self.mapping.moe.ep_rank, + ep_size=self.mapping.moe.ep_size, + ) + for raw_name, loaded_weight in weights: + name = self._map_checkpoint_name(raw_name) + if name is None: + continue + + if ".indexer." in name: + if name.endswith(".bias") and name not in params_dict: + continue + param = self.get_param(params_dict, name) + if param is not None: + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + self._try_load_fused_indexer_projection( + name=name, + loaded_weight=loaded_weight, + params_dict=params_dict, + modules_dict=modules_dict, + pending_fp8_wk=pending_fp8_wk, + loaded_shards=loaded_fused_indexer_shards, + ) + continue + + for param_name, weight_name, shard_id in _STACKED_PARAMS_MAPPING: + if weight_name not in name: + continue + if ("mlp.experts." in name) and name not in params_dict: + continue + name = name.replace(weight_name, param_name) + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + if name.endswith(".bias") and name not in params_dict: + continue + if moe_loader.matches(name): + moe_loader.load(name, loaded_weight) + continue + + if cached_a_proj is not None and ( + "q_a_proj" in name or "kv_a_proj_with_mqa" in name + ): + cached_a_proj[name] = loaded_weight + q_a_proj_name = ( + name + if "q_a_proj" in name + else name.replace("kv_a_proj_with_mqa", "q_a_proj") + ) + kv_a_proj_name = ( + name + if "kv_a_proj_with_mqa" in name + else name.replace("q_a_proj", "kv_a_proj_with_mqa") + ) + + if ( + q_a_proj_name in cached_a_proj + and kv_a_proj_name in cached_a_proj + ): + q_a_proj_weight = cached_a_proj[q_a_proj_name] + kv_a_proj_weight = cached_a_proj[kv_a_proj_name] + fused_weight = torch.cat( + [q_a_proj_weight, kv_a_proj_weight], dim=0 + ) + + if "q_a_proj" in name: + param_name = name.replace( + "q_a_proj", "fused_qkv_a_proj_with_mqa" + ) + else: + param_name = name.replace( + "kv_a_proj_with_mqa", "fused_qkv_a_proj_with_mqa" + ) + param = params_dict[param_name] + + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, fused_weight) + cached_a_proj.pop(q_a_proj_name) + cached_a_proj.pop(kv_a_proj_name) + else: + if ".mlp.experts." in name: + continue + param = self.get_param(params_dict, name) + if param is None: + continue + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + self.post_load_weights() + + def post_load_weights(self) -> None: + self_attn = self.model.decoder.self_attn + pad_fused_qkv_a_proj_weight_for_fp8_blockscale(self_attn) + if ( + hasattr(self.quant_config, "weight_block_size") + and (self.quant_config.weight_block_size is not None) + and self_attn.kv_b_proj.weight.dtype + in ( + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + ) + ): + weight_block_size = self.quant_config.weight_block_size + dtype = torch.get_default_dtype() + w = block_dequant( + self_attn.kv_b_proj.weight, + self_attn.kv_b_proj.weight_scale_inv, + weight_block_size, + ).to(dtype) + else: + w = self_attn.kv_b_proj.weight + + 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) + self_attn.w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2) + self_attn.w_vc = w_vc.contiguous().transpose(1, 2) + + +EntryClass = [GlmMoeDsaForCausalLMNextN] diff --git a/python/tokenspeed/runtime/models/gpt_oss.py b/python/tokenspeed/runtime/models/gpt_oss.py new file mode 100644 index 0000000..678db6f --- /dev/null +++ b/python/tokenspeed/runtime/models/gpt_oss.py @@ -0,0 +1,994 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Inference-only GptOss model compatible with HuggingFace weights.""" + +# ruff: noqa: E402 + +from __future__ import annotations + +import math +import re +from collections.abc import Iterable +from typing import Any + +import torch +from torch import nn +from transformers import PretrainedConfig + +from tokenspeed.runtime.configs.utils import get_rope_theta +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.distributed.process_group_manager import ( + process_group_manager as pg_manager, +) +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.layers.linear import ( + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from tokenspeed.runtime.layers.moe import ( + ExpertCheckpointSchema, + build_moe_checkpoint_loader, +) +from tokenspeed.runtime.layers.moe.expert import MoELayer +from tokenspeed.runtime.layers.moe.topk import TopK +from tokenspeed.runtime.layers.moe.utils import get_all2all_backend +from tokenspeed.runtime.layers.paged_attention import PagedAttention +from tokenspeed.runtime.layers.quantization import QuantizationConfig +from tokenspeed.runtime.layers.rotary_embedding import get_rope +from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader +from tokenspeed.runtime.models.base import ( + BaseCausalLM, + BaseTransformerModel, + CompiledMoEDecoderLayer, +) +from tokenspeed.runtime.models.utils import ( + create_fused_set_kv_buffer_arg, + validate_attention_partition, +) +from tokenspeed.runtime.utils import add_prefix, get_colorful_logger +from tokenspeed.runtime.utils.env import global_server_args_dict +from tokenspeed.runtime.utils.pdl import pdl_enabled + +logger = get_colorful_logger(__name__) + + +from tokenspeed_kernel.ops.gemm.flashinfer import tinygemm_bf16 +from tokenspeed_kernel.registry import error_fn + + +class TinyGemmLinear(ReplicatedLinear): + """ReplicatedLinear with a FlashInfer tinygemm BF16 fast path for small batches.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._use_tinygemm = ( + tinygemm_bf16 is not error_fn + and not self.skip_bias_add + and self.weight.is_contiguous() + and self.weight.shape[0] % 16 == 0 + and self.weight.shape[1] % 64 == 0 + and self.weight.dtype == torch.bfloat16 + and ( + self.bias is None + or ( + self.bias.dtype == torch.bfloat16 + and self.bias.is_contiguous() + and self.bias.shape[0] == self.weight.shape[0] + ) + ) + ) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]: + if ( + self._use_tinygemm + and x.ndim == 2 + and x.is_cuda + and x.shape[0] <= 128 + and x.is_contiguous() + and x.shape[1] == self.weight.shape[1] + and x.dtype == torch.bfloat16 + ): + out = x.new_empty((x.shape[0], self.output_size)) + tinygemm_bf16(x, self.weight, out, self.bias, use_pdl=pdl_enabled()) + return out, None + + return super().forward(x) + + +class GptOssAttention(nn.Module): + def __init__( + self, + config, + mapping: Mapping, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + layer_id: int = 0, + rope_theta: float = 10000, + rope_scaling: dict[str, Any] | None = None, + max_position_embeddings: int = 8192, + head_dim: int | None = None, + rms_norm_eps: float = 1e-06, + attention_bias: bool = False, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + sliding_window_size: int = -1, + layer_type: str = "", + params_dtype: torch.dtype = torch.bfloat16, + ) -> None: + + super().__init__() + self.mapping = mapping + self.hidden_size = hidden_size + self.sliding_window_size = sliding_window_size + + attn_tp_rank = self.mapping.attn.tp_rank + attn_tp_size = self.mapping.attn.tp_size + attn_tp_group = self.mapping.attn.tp_group + + self.total_num_heads = num_heads + self.total_num_kv_heads = num_kv_heads + validate_attention_partition( + self.total_num_heads, + self.total_num_kv_heads, + attn_tp_size, + ) + self.num_heads = self.total_num_heads // attn_tp_size + self.num_kv_heads = max(1, self.total_num_kv_heads // attn_tp_size) + self.head_dim = head_dim or hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + self.tp_rank = self.mapping.rank + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=attention_bias, + params_dtype=params_dtype, + quant_config=quant_config, + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + tp_group=attn_tp_group, + prefix=add_prefix("qkv_proj", prefix), + ) + + self.sinks = nn.Parameter( + torch.empty(self.num_heads, dtype=torch.bfloat16), requires_grad=False + ) + + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=attention_bias, + quant_config=quant_config, + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + tp_group=attn_tp_group, + reduce_results=False, + params_dtype=params_dtype, + prefix=add_prefix("o_proj", prefix), + ) + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + ) + + if layer_type not in {"sliding_attention", "full_attention"}: + raise ValueError(f"Unsupported attention layer_type: {layer_type}.") + use_sliding_window = layer_type == "sliding_attention" + self.attn = PagedAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + sliding_window_size=(sliding_window_size if use_sliding_window else -1), + group_id=layer_type, + ) + self.layer_id = layer_id + + def forward_prepare( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + ): + + if hidden_states.shape[0] == 0: + return hidden_states, ctx, out_cache_loc, None + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + fused_kv_arg = None + if ctx.attn_backend.support_kv_cache_prewrite(ctx.forward_mode): + n = q.shape[0] + v_3d = v.view(n, self.num_kv_heads, self.head_dim) + fused_kv_arg = create_fused_set_kv_buffer_arg( + value=v_3d, + layer=self.attn, + # Flat path: prewrite at this layer's group locations. + out_cache_loc=ctx.attn_backend.select_out_cache_loc( + self.attn, out_cache_loc, ctx.forward_mode + ), + token_to_kv_pool=ctx.token_to_kv_pool, + ) + + if fused_kv_arg is not None: + n = q.shape[0] + q_rope = torch.empty((n, self.q_size), dtype=q.dtype, device=q.device) + q, k = self.rotary_emb( + positions, + q, + k, + fused_set_kv_buffer_arg=fused_kv_arg, + output_q_rope=q_rope, + enable_pdl=pdl_enabled(), + ) + inner_state = q_rope, None, None + else: + q, k = self.rotary_emb(positions, q, k) + inner_state = q, k, v + return None, ctx, out_cache_loc, inner_state + + def forward_core(self, intermediate_state): + + hidden_states, ctx, out_cache_loc, inner_state = intermediate_state + if inner_state is None: + return hidden_states + # Cache was already written by the fused RoPE+KV kernel iff we took that path, + # which is exactly when k is None in inner_state. + save_kv_cache = inner_state[1] is not None + attn_output = self.attn( + *inner_state, + save_kv_cache=save_kv_cache, + ctx=ctx, + out_cache_loc=out_cache_loc, + sinks=self.sinks, + ) + output, _ = self.o_proj(attn_output) + return output + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + ) -> torch.Tensor: + + s = self.forward_prepare( + positions=positions, + hidden_states=hidden_states, + ctx=ctx, + out_cache_loc=out_cache_loc, + ) + return self.forward_core(s) + + +def routing_function(hidden_states, gating_output, topk, renormalize): + + experts = torch.topk(gating_output, k=topk, dim=-1, sorted=True) + expert_weights = torch.nn.functional.softmax( + experts.values.to(torch.float32), dim=1 + ) + expert_indices = experts.indices.to(torch.int32) + return expert_weights, expert_indices + + +class GptOssSparseMoeBlock(nn.Module): + def __init__( + self, + config, + mapping: Mapping, + num_experts: int, + top_k: int, + hidden_size: int, + intermediate_size: int, + params_dtype: torch.dtype | None = None, + quant_config: QuantizationConfig | None = None, + layer_index: int = -1, + prefix: str = "", + ): + + super().__init__() + self.mapping = mapping + self.layer_index = layer_index + self.tp_size = self.mapping.world_size + self.hidden_size = hidden_size + self.activation = config.hidden_act + self.activation_alpha = getattr(config, "hidden_act_alpha", 1.702) + self.swiglu_limit = config.swiglu_limit + self.num_experts = ( + num_experts + global_server_args_dict["ep_num_redundant_experts"] + ) + self.quant_config = quant_config + if self.tp_size > config.num_local_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than " + f"the number of experts {config.num_local_experts}." + ) + + self.experts = MoELayer( + top_k=top_k, + num_experts=self.num_experts, + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + quant_config=self.quant_config, + layer_index=self.layer_index, + prefix=add_prefix("experts", prefix), + tp_rank=self.mapping.moe.tp_rank, + tp_size=self.mapping.moe.tp_size, + ep_rank=self.mapping.moe.ep_rank, + ep_size=self.mapping.moe.ep_size, + activation="swiglu", + activation_alpha=self.activation_alpha, + swiglu_limit=self.swiglu_limit, + # HF gpt-oss stores ``gate_up_proj_blocks`` row-interleaved + # ([w1_0, w3_0, w1_1, w3_1, ...]) and uses the gpt-oss SwiGLU+1 + # activation silu(α·gate)·(up + 1). + swiglu_beta=1.0, + w13_input_layout="interleaved", + with_bias=True, + ) + + self.router = TinyGemmLinear( + config.hidden_size, + config.num_local_experts, + bias=True, + quant_config=None, + prefix=add_prefix("gate", prefix), + params_dtype=config.dtype, + ) + + self.topk = TopK( + top_k=top_k, + custom_routing_function=routing_function, + output_format=self.experts.topk_output_format, + topk_indices_dtype=( + torch.int64 if get_all2all_backend().is_deepep() else torch.int32 + ), + ) + + def forward( + self, + hidden_states: torch.Tensor, + num_global_tokens: int, + max_num_tokens_per_gpu: int, + ) -> torch.Tensor: + + # router_logits: (num_tokens, n_experts) + if hidden_states.shape[0] == 0: + router_logits = hidden_states.new_empty(0, self.router.weight.shape[0]) + else: + router_output = self.router(hidden_states) + router_logits = ( + router_output[0] if isinstance(router_output, tuple) else router_output + ) + if hidden_states.shape[0] > 0: + topk_output = self.topk(hidden_states, router_logits) + else: + topk_output = self.topk.empty_topk_output( + hidden_states.device, + hidden_states=hidden_states, + router_logits=router_logits, + ) + return self.experts( + hidden_states=hidden_states, + topk_output=topk_output, + num_global_tokens=num_global_tokens, + max_num_tokens_per_gpu=max_num_tokens_per_gpu, + ) + + def get_moe_weights(self) -> list[torch.Tensor]: + + return [ + x.data + for name, x in self.experts.named_parameters() + if name not in ["correction_bias"] + ] + + +class _WeightCreator: + def __init__(self, fn): + self._fn = fn + + @staticmethod + def maybe_materialize(obj): + if isinstance(obj, _WeightCreator): + output = obj._fn() + obj._fn = None + return output + return obj + + +class GptOssConfig(PretrainedConfig): + model_type = "gpt_oss" + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + +def get_attention_sliding_window_size(config): + # Aligned with HF's implementation, using sliding window inclusive with the last token + # TokenSpeed assumes exclusive + return config.sliding_window - 1 + + +class GptOssDecoderLayer(CompiledMoEDecoderLayer): + + def __init__( + self, + config: GptOssConfig, + layer_id: int, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + sliding_window_size: int | None = None, + ) -> None: + + self._config = config + self._mapping = mapping + self._quant_config = quant_config + self._prefix = prefix + + if sliding_window_size is None: + self.sliding_window_size = get_attention_sliding_window_size(config) + else: + self.sliding_window_size = sliding_window_size + + super().__init__( + config=config, + layer_id=layer_id, + mapping=mapping, + quant_config=quant_config, + prefix=prefix, + ) + + self.attn_tp_group = pg_manager.get_process_group( + "nccl", self.mapping.attn.tp_group + ) + self.attn_tp_size = self.mapping.attn.tp_size + self.attn_tp_rank = self.mapping.attn.tp_rank + + def resolve_attn(self, prefix: str) -> nn.Module: + + config = self._config + head_dim = getattr( + config, "head_dim", config.hidden_size // config.num_attention_heads + ) + + return GptOssAttention( + config=config, + mapping=self._mapping, + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + layer_id=self.layer_id, + rope_theta=get_rope_theta(config), + rope_scaling=getattr(config, "rope_scaling", None), + max_position_embeddings=getattr(config, "max_position_embeddings", 8192), + head_dim=head_dim, + rms_norm_eps=config.rms_norm_eps, + attention_bias=config.attention_bias, + quant_config=self._quant_config, + prefix=add_prefix("self_attn", prefix), + sliding_window_size=self.sliding_window_size, + layer_type=config.layer_types[self.layer_id], + params_dtype=config.dtype, + ) + + def resolve_mlp(self, prefix: str) -> nn.Module: + + config = self._config + + return GptOssSparseMoeBlock( + config=config, + mapping=self._mapping, + num_experts=config.num_local_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + quant_config=self._quant_config, + layer_index=self.layer_id, + prefix=add_prefix("mlp", prefix), + ) + + +class GptOssModel(BaseTransformerModel): + layer_cls = GptOssDecoderLayer + + +class GptOssForCausalLM(BaseCausalLM): + model_cls = GptOssModel + fall_back_to_pt_during_load = False + + def get_attention_sliding_window_size(self): + return get_attention_sliding_window_size(self.config) + + @classmethod + def get_model_config_for_expert_location(cls, config): + from tokenspeed.runtime.moe.expert_location import ( + ModelConfigForExpertLocation, + ) + + return ModelConfigForExpertLocation( + num_layers=config.num_hidden_layers, + num_logical_experts=config.num_local_experts, + num_groups=None, + ) + + def _get_default_weight_mapping(self): + + weight_mapping = {} + weight_mapping["embedding.weight"] = "model.embed_tokens.weight" + weight_mapping["unembedding.weight"] = "lm_head.weight" + weight_mapping["norm.scale"] = "model.norm.weight" + + for layer_id in range(self.config.num_hidden_layers): + pfx = f"model.layers.{layer_id}" + bpfx = f"block.{layer_id}" + + for proj in ("q_proj", "k_proj", "v_proj"): + weight_mapping[f"{bpfx}.attn.{proj}.weight"] = ( + f"{pfx}.self_attn.{proj}.weight" + ) + weight_mapping[f"{bpfx}.attn.{proj}.bias"] = ( + f"{pfx}.self_attn.{proj}.bias" + ) + + weight_mapping[f"{bpfx}.attn.out.weight"] = f"{pfx}.self_attn.o_proj.weight" + weight_mapping[f"{bpfx}.attn.out.bias"] = f"{pfx}.self_attn.o_proj.bias" + weight_mapping[f"{bpfx}.attn.sinks"] = f"{pfx}.self_attn.sinks" + weight_mapping[f"{bpfx}.attn.norm.scale"] = f"{pfx}.input_layernorm.weight" + + weight_mapping[f"{bpfx}.mlp.gate.weight"] = f"{pfx}.mlp.router.weight" + weight_mapping[f"{bpfx}.mlp.gate.bias"] = f"{pfx}.mlp.router.bias" + weight_mapping[f"{bpfx}.mlp.norm.scale"] = ( + f"{pfx}.post_attention_layernorm.weight" + ) + weight_mapping[f"{bpfx}.mlp.experts.gate_up_proj"] = ( + f"{pfx}.mlp.experts.gate_up_proj" + ) + weight_mapping[f"{bpfx}.mlp.gate_up_proj_bias"] = ( + f"{pfx}.mlp.experts.gate_up_proj_bias" + ) + weight_mapping[f"{bpfx}.mlp.down_proj"] = f"{pfx}.mlp.experts.mlp2_weight" + weight_mapping[f"{bpfx}.mlp.down_proj_bias"] = ( + f"{pfx}.mlp.experts.mlp2_bias" + ) + + return weight_mapping + + def load_weights( + self, + weights: Iterable[tuple[str, torch.Tensor]], + is_nextn: bool = False, + weight_name_mapping: dict = None, + ): + + quant_config_name = ( + self.quant_config.get_name() if self.quant_config is not None else None + ) + if is_nextn: + raise ValueError("GPT-OSS does not support nextn weight loading.") + + if quant_config_name == "mxfp4": + self._load_mxfp4_weights(weights, weight_name_mapping=weight_name_mapping) + else: + self._load_normal_weights(weights, weight_name_mapping=weight_name_mapping) + + def _load_normal_weights( + self, + weights, + weight_name_mapping: dict = None, + other_loaded_param_names: set = None, + ): + + attn_tp_rank = self.mapping.attn.tp_rank + rank = self.mapping.rank + weights = sorted(weights, key=lambda x: x[0]) + + if weight_name_mapping is None: + weight_name_mapping = self._get_default_weight_mapping() + else: + default_mapping = self._get_default_weight_mapping() + default_mapping.update(weight_name_mapping) + weight_name_mapping = default_mapping + + stacked_params_mapping = [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ] + + params_dict = dict(self.named_parameters()) + # MoE expert weights, scales, and activation scales are handled + # by the checkpoint loader. + moe_loader = build_moe_checkpoint_loader( + params_dict=params_dict, + fused_schema=ExpertCheckpointSchema( + gate_up_fused_name="gate_up_proj", + down_proj_name="down_proj", + extra_names={ + "gate_up_bias": "gate_up_proj_bias", + "down_bias": "down_proj_bias", + }, + ), + num_experts=self.config.num_local_experts, + ep_rank=self.mapping.moe.ep_rank, + ep_size=self.mapping.moe.ep_size, + fused_gate_up_as_w13=True, + include_bias=True, + fused_load_style="local_tensor", + transpose_local_tensor_non_bias=True, + ) + params_checker = {k: False for k in params_dict} + + for name, loaded_weight in weights: + loaded_weight = _WeightCreator.maybe_materialize(loaded_weight) + + if weight_name_mapping and name in weight_name_mapping: + name = weight_name_mapping[name] + + if "rotary_emb.inv_freq" in name: + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if "mlp.experts" in name: + continue + + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if name not in params_dict: + continue + + param = params_dict[name] + param.weight_loader(param, loaded_weight, shard_id) + params_checker[name] = True + break + + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if moe_loader.matches(name): + mapped_name = moe_loader.load(name, loaded_weight) + params_checker[mapped_name] = True + name = mapped_name + else: + if name not in params_dict: + continue + param = params_dict[name] + if "sinks" in name: + start = attn_tp_rank * param.numel() + param.data.copy_(loaded_weight[start : start + param.numel()]) + else: + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + params_checker[name] = True + + not_loaded_params = [] + already_loaded = other_loaded_param_names or set() + for k, v in params_checker.items(): + if ( + not v + and ("weight_scale" not in k) + and ("input_scale" not in k) + and k not in already_loaded + ): + not_loaded_params.append(k) + + if rank == 0: + if len(not_loaded_params) > 0: + raise RuntimeError(f"Not all parameters loaded: {not_loaded_params=}") + else: + logger.info("All parameters loaded successfully.") + + self.routed_experts_weights_of_layer = { + layer_id: self.model.layers[layer_id].mlp.get_moe_weights() + for layer_id in range(len(self.model.layers)) + } + + def _load_mxfp4_weights(self, weights, weight_name_mapping: dict): + + mxfp4_weights = [] + normal_weights = [] + + for name, weight in weights: + if ".experts" in name: + mxfp4_weights.append((name, weight)) + else: + normal_weights.append((name, weight)) + + mxfp4_loaded_params = self._load_mxfp4_experts_weights(mxfp4_weights) + self._load_normal_weights( + normal_weights, + weight_name_mapping=weight_name_mapping, + other_loaded_param_names=mxfp4_loaded_params, + ) + + def _load_mxfp4_experts_weights(self, weights): + + params_dict = dict(self.named_parameters()) + loaded_params: set = set() + mxfp4_block = 32 + + moe_tp_rank = self.mapping.moe.tp_rank + moe_tp_size = self.mapping.moe.tp_size + moe_ep_rank = self.mapping.moe.ep_rank + moe_ep_size = self.mapping.moe.ep_size + + intermediate_size = self.config.intermediate_size + intermediate_size_block = intermediate_size // mxfp4_block + per_rank_intermediate_size_block = math.ceil( + intermediate_size_block / moe_tp_size + ) + per_rank_intermediate_size = per_rank_intermediate_size_block * mxfp4_block + + moe_num_global_experts = self.config.num_local_experts + moe_num_local_experts = moe_num_global_experts // moe_ep_size + + moe_tp_rank_start = moe_tp_rank * per_rank_intermediate_size + moe_tp_rank_end = min( + (moe_tp_rank + 1) * per_rank_intermediate_size, intermediate_size + ) + + moe_ep_rank_start = moe_ep_rank * moe_num_local_experts + moe_ep_rank_end = (moe_ep_rank + 1) * moe_num_local_experts + + def _copy_into_param(param, narrow_weight): + if param.shape == narrow_weight.shape: + param.data.copy_(narrow_weight) + else: + slices = tuple( + slice(0, min(p, n)) + for p, n in zip(param.shape, narrow_weight.shape) + ) + param.data[slices].copy_(narrow_weight[slices]) + + # Detect AMD-Quark per-expert checkpoints (e.g. + # ``amd/gpt-oss-120b-w-mxfp4-a-fp8``). These store one set of tensors + # per expert (``...experts.{e}.gate_up_proj.{weight,...}``) plus a + # scalar ``input_scale`` for static FP8 activation quantization. + if any( + re.search(r"\.experts\.\d+\.(gate_up_proj|down_proj)\.", n) + for n, _ in weights + ): + return self._load_mxfp4_per_expert_weights( + weights, + params_dict=params_dict, + moe_tp_rank_start=moe_tp_rank_start, + moe_tp_rank_end=moe_tp_rank_end, + moe_ep_rank_start=moe_ep_rank_start, + moe_ep_rank_end=moe_ep_rank_end, + moe_tp_rank=moe_tp_rank, + copy_into_param=_copy_into_param, + mxfp4_block=mxfp4_block, + ) + + for name, weight in weights: + weight = _WeightCreator.maybe_materialize(weight) + + if "gate_up_proj_blocks" in name: + new_name = name.replace("gate_up_proj_blocks", "w13_weight") + weight = weight.view( + moe_num_global_experts, 2 * intermediate_size, -1 + ).contiguous() + narrow_weight = weight[ + moe_ep_rank_start:moe_ep_rank_end, + 2 * moe_tp_rank_start : 2 * moe_tp_rank_end, + ..., + ] + _copy_into_param(params_dict[new_name], narrow_weight) + loaded_params.add(new_name) + + elif "down_proj_blocks" in name: + new_name = name.replace("down_proj_blocks", "w2_weight") + weight = weight.view( + moe_num_global_experts, -1, intermediate_size // 2 + ).contiguous() + narrow_weight = weight[ + moe_ep_rank_start:moe_ep_rank_end, + ..., + moe_tp_rank_start // 2 : moe_tp_rank_end // 2, + ] + _copy_into_param(params_dict[new_name], narrow_weight) + loaded_params.add(new_name) + + elif "gate_up_proj_scales" in name: + new_name = name.replace("gate_up_proj_scales", "w13_weight_scale") + narrow_weight = weight[ + moe_ep_rank_start:moe_ep_rank_end, + 2 * moe_tp_rank_start : 2 * moe_tp_rank_end, + ..., + ] + _copy_into_param(params_dict[new_name], narrow_weight) + loaded_params.add(new_name) + + elif "down_proj_scales" in name: + new_name = name.replace("down_proj_scales", "w2_weight_scale") + narrow_weight = weight[ + moe_ep_rank_start:moe_ep_rank_end, + ..., + moe_tp_rank_start // mxfp4_block : moe_tp_rank_end // mxfp4_block, + ] + _copy_into_param(params_dict[new_name], narrow_weight) + loaded_params.add(new_name) + + elif "gate_up_proj_bias" in name: + new_name = name.replace("gate_up_proj_bias", "w13_weight_bias") + narrow_weight = weight[ + moe_ep_rank_start:moe_ep_rank_end, + 2 * moe_tp_rank_start : 2 * moe_tp_rank_end, + ] + _copy_into_param(params_dict[new_name], narrow_weight) + loaded_params.add(new_name) + + elif "down_proj_bias" in name: + new_name = name.replace("down_proj_bias", "w2_weight_bias") + narrow_weight = weight[moe_ep_rank_start:moe_ep_rank_end, ...] + if moe_tp_rank != 0: + narrow_weight = torch.zeros_like(narrow_weight) + _copy_into_param(params_dict[new_name], narrow_weight) + loaded_params.add(new_name) + + return loaded_params + + def _load_mxfp4_per_expert_weights( + self, + weights, + *, + params_dict, + moe_tp_rank_start: int, + moe_tp_rank_end: int, + moe_ep_rank_start: int, + moe_ep_rank_end: int, + moe_tp_rank: int, + copy_into_param, + mxfp4_block: int, + ): + """Load the AMD-Quark per-expert MXFP4 + FP8 input-scale checkpoint. + + Tensor names look like + ``model.layers.{l}.mlp.experts.{e}.{gate_up_proj,down_proj}.{weight, + weight_scale,bias,input_scale}`` and shapes match the existing fused + ``w13_*`` / ``w2_*`` parameters once the per-expert tensors are + stacked along the expert dimension. + """ + loaded_params: set = set() + + per_expert_re = re.compile( + r"^(?P.*\.experts\.)(?P\d+)\.(?Pgate_up_proj|down_proj)\.(?Pweight_scale|weight|bias|input_scale)$" + ) + + for name, weight in weights: + weight = _WeightCreator.maybe_materialize(weight) + + match = per_expert_re.match(name) + if match is None: + # ``router`` and other non-expert weights are emitted to the + # generic loader by the caller; if we still hit one here it is + # an unexpected name. + continue + + base = match.group("base") + expert_id = int(match.group("expert")) + proj = match.group("proj") + kind = match.group("kind") + + if not (moe_ep_rank_start <= expert_id < moe_ep_rank_end): + continue + local_expert_id = expert_id - moe_ep_rank_start + + if proj == "gate_up_proj": + if kind == "weight": + target = base + "w13_weight" + elif kind == "weight_scale": + target = base + "w13_weight_scale" + elif kind == "bias": + target = base + "w13_weight_bias" + else: # input_scale + target = base + "w13_input_scale" + else: # down_proj + if kind == "weight": + target = base + "w2_weight" + elif kind == "weight_scale": + target = base + "w2_weight_scale" + elif kind == "bias": + target = base + "w2_weight_bias" + else: # input_scale + target = base + "w2_input_scale" + + if target not in params_dict: + # The active backend (e.g. plain MXFP4 without FP8 act) may + # not allocate ``input_scale`` parameters; just skip. + if kind == "input_scale": + continue + raise KeyError(f"missing target parameter {target!r} for {name!r}") + param = params_dict[target] + + if kind == "input_scale": + # Per-tensor static FP8 activation scale; broadcast scalar + # into the per-expert slot. + param.data[local_expert_id] = ( + weight.detach().to(torch.float32).reshape(()) + ) + loaded_params.add(target) + continue + + if proj == "gate_up_proj": + # Per-expert tensor shapes: + # weight: (2*intermediate, hidden//2) uint8 + # weight_scale: (2*intermediate, hidden//mxfp4_block) uint8 + # bias: (2*intermediate,) bf16 + # The fused parameter slot is sharded along the (output) + # intermediate dimension. + if kind == "bias": + narrow = weight[2 * moe_tp_rank_start : 2 * moe_tp_rank_end] + else: + narrow = weight[2 * moe_tp_rank_start : 2 * moe_tp_rank_end, :] + copy_into_param(param.data[local_expert_id], narrow) + else: # down_proj + # Per-expert tensor shapes: + # weight: (hidden, intermediate//2) uint8 + # weight_scale: (hidden, intermediate//mxfp4_block) uint8 + # bias: (hidden,) bf16 + # Down_proj is sharded along the (input) intermediate + # dimension. + if kind == "bias": + if moe_tp_rank != 0: + narrow = torch.zeros_like(weight) + else: + narrow = weight + elif kind == "weight": + narrow = weight[:, moe_tp_rank_start // 2 : moe_tp_rank_end // 2] + else: # weight_scale + narrow = weight[ + :, + moe_tp_rank_start + // mxfp4_block : moe_tp_rank_end + // mxfp4_block, + ] + copy_into_param(param.data[local_expert_id], narrow) + + loaded_params.add(target) + + return loaded_params + + +EntryClass = GptOssForCausalLM diff --git a/python/tokenspeed/runtime/models/kimi_k25.py b/python/tokenspeed/runtime/models/kimi_k25.py new file mode 100644 index 0000000..afecf47 --- /dev/null +++ b/python/tokenspeed/runtime/models/kimi_k25.py @@ -0,0 +1,1056 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright 2023-2024 SGLang Team +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Inference-only Kimi-K2.5 VLM (DeepseekV3 LM + MoonViT vision tower) compatible with HuggingFace weights.""" + +from __future__ import annotations + +import logging +import math +from collections.abc import Iterable, Sequence +from copy import deepcopy + +import numpy as np +import torch +import torch.nn.functional as F +from torch import nn +from transformers import activations + +from tokenspeed.runtime.configs.kimi_k25_config import ( + KimiK25Config, + KimiK25VisionConfig, +) +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.layers.conv import Conv2dLayer +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.moe.expert_location import ModelConfigForExpertLocation +from tokenspeed.runtime.multimodal.embedder import ( + EncoderSpec, + VisionEmbedder, + pad_input_tokens, +) + +try: + from transformers.activations import PytorchGELUTanh +except ImportError: + from transformers.activations import GELUTanh + + activations.PytorchGELUTanh = GELUTanh + PytorchGELUTanh = GELUTanh + +from tokenspeed.runtime.layers.attention.mm_encoder_attention import VisionAttention +from tokenspeed.runtime.layers.linear import ReplicatedLinear + +try: + from tokenspeed.runtime.layers.quantization.modelslim.modelslim import ( + ModelSlimConfig, + ) +except ImportError: + + class ModelSlimConfig: + pass + + +try: + from tokenspeed.runtime.layers.quantization.quark.quark import QuarkConfig +except ImportError: + + class QuarkConfig: + pass + + +from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader +from tokenspeed.runtime.models.deepseek_v3 import DeepseekV3ForCausalLM +from tokenspeed.runtime.multimodal.encoder_cudagraph import ( + EncoderCudaGraphWrapper, + VisionEncoderCudaGraphAdapter, +) +from tokenspeed.runtime.multimodal.inputs import ( + Modality, + MultimodalDataItem, + MultimodalInputs, +) +from tokenspeed.runtime.utils import add_prefix + +logger = logging.getLogger(__name__) + + +class MLP2(nn.Module): + """ + Two-layer MLP helper used by the Kimi-K2.5 MoonViT blocks. + + This helper is inlined so the TokenSpeed VLM snapshot can keep only the + Kimi-K2.5 target model. + """ + + def __init__( + self, + dims: list[int], + activation, + bias: bool = True, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + if len(dims) != 3: + raise ValueError(f"dims must have length 3, got {len(dims)}.") + + self.quant_config = quant_config + if isinstance(self.quant_config, ModelSlimConfig): + self.fc0 = ReplicatedLinear( + dims[0], + dims[1], + bias=bias, + quant_config=quant_config, + prefix=add_prefix("fc0", prefix), + ) + self.fc1 = ReplicatedLinear( + dims[1], + dims[2], + bias=bias, + quant_config=quant_config, + prefix=add_prefix("fc1", prefix), + ) + else: + self.fc0 = nn.Linear(dims[0], dims[1], bias=bias) + self.fc1 = nn.Linear(dims[1], dims[2], bias=bias) + for module in (self.fc0, self.fc1): + nn.init.trunc_normal_( + module.weight, std=math.sqrt(2 / module.in_features) + ) + if module.bias is not None: + nn.init.zeros_(module.bias) + self.activation = activation + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if isinstance(self.quant_config, ModelSlimConfig): + x = x.flatten(0, 1) + x, _ = self.fc0(x) + x = self.activation(x) + x, _ = self.fc1(x) + else: + x = self.fc0(x) + x = self.activation(x) + x = self.fc1(x) + return x + + +def apply_rope( + xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor, x_shape=None +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Args: (The leading dimensions of all inputs should be the same) + xq: query, tensor of shape (..., num_heads, head_dim) + xk: key, tensor of shape (..., num_heads, head_dim) + freqs_cis: tensor of shape (..., head_dim/2), dtype=torch.complex64. It contains the precomputed cis(freqs) for each position in the 2D grid. + Returns: + xq_out, xk_out: tensors of shape (..., num_heads, head_dim) + """ + + freqs_cis = freqs_cis.unsqueeze(-2) # ..., 1, head_dim/2 + # ..., num_heads, head_dim/2 + xq_ = torch.view_as_complex(xq.float().view(*xq.shape[:-1], -1, 2)) + xk_ = torch.view_as_complex(xk.float().view(*xq.shape[:-1], -1, 2)) + xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(-2) # ..., num_heads, head_dim + xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(-2) # ..., num_heads, head_dim + return xq_out.type_as(xq), xk_out.type_as(xk) + + +def tpool_patch_merger( + x: torch.Tensor, + grid_thws: torch.Tensor, + merge_kernel_size: tuple[int, int] = (2, 2), +) -> list[torch.Tensor]: + d_model = x.size(-1) + + outputs = [] + pre_sum = 0 + for t, h, w in grid_thws.tolist(): + seq = x[pre_sum : pre_sum + t * h * w] + kernel_height, kernel_width = merge_kernel_size + new_height, new_width = h // kernel_height, w // kernel_width + reshaped_seq = seq.view( + t, new_height, kernel_height, new_width, kernel_width, d_model + ) + reshaped_seq = ( + reshaped_seq.permute(0, 1, 3, 2, 4, 5).contiguous().mean(dim=0) + ) # temporal pooling + padded_seq = reshaped_seq.view( + new_height * new_width, kernel_height * kernel_width, -1 + ) + outputs.append(padded_seq) + pre_sum += t * h * w + + return outputs + + +class MoonViTEncoderLayer(nn.Module): + + def __init__( + self, + num_heads: int, + hidden_dim: int, + mlp_dim: int, + mapping: Mapping, + *, + activation=F.gelu, + attn_bias: bool = False, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + mm_attention_backend: str | None = None, + ): + super().__init__() + self.num_heads = num_heads + self.hidden_dim = hidden_dim + + self.norm0 = nn.LayerNorm(hidden_dim) + self.norm1 = nn.LayerNorm(hidden_dim) + + self.mlp = MLP2( + [hidden_dim, mlp_dim, hidden_dim], + activation, + quant_config=quant_config, + prefix=add_prefix("mlp", prefix), + ) + + self.attn = VisionAttention( + embed_dim=hidden_dim, + num_heads=num_heads, + qkv_bias=attn_bias, + proj_bias=attn_bias, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + customized_position_embedding_applier=apply_rope, + position_embedding_mode="complex_rope", + mapping=mapping, + mm_attention_backend=mm_attention_backend, + ) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: int, + rope_freqs_cis: torch.Tensor | None = None, + ): + if not isinstance(max_seqlen, int): + raise TypeError( + f"max_seqlen must be a Python int for capture-safety, got {type(max_seqlen)}" + ) + residual = hidden_states + hidden_states = self.norm0(hidden_states) + + hidden_states = self.attn( + hidden_states, + cu_seqlens=cu_seqlens, + position_embeddings=rope_freqs_cis, + max_seqlen=max_seqlen, + ) + + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.norm1(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +def get_rope_shape_decorate(func): + _get_rope_shape_first_call_flag = set() + + def wrapper(org, interpolation_mode, shape): + key = (org.requires_grad, torch.is_grad_enabled(), interpolation_mode) + if key not in _get_rope_shape_first_call_flag: + _get_rope_shape_first_call_flag.add(key) + _ = func(org, interpolation_mode, shape=(64, 64)) + return func(org, interpolation_mode, shape) + + return wrapper + + +def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): + """ + From: + https://github.com/OpenGVLab/InternVideo/blob/421f6d2361fc8f61a3394244571f2601a4e99e29/InternVideo2/multi_modality/models/backbones/internvideo2/pos_embed.py#L86 + embed_dim: output dimension for each position + pos: a list of positions to be encoded: size (M,) + out: (M, D) + """ + if embed_dim % 2 != 0: + raise ValueError(f"embed_dim must be even, got {embed_dim}.") + omega = np.arange(embed_dim // 2, dtype=np.float32) + omega /= embed_dim / 2.0 + omega = 1.0 / 10000**omega # (D/2,) + + pos = pos.reshape(-1) # (M,) + out = np.einsum("m,d->md", pos, omega) # (M, D/2), outer product + + emb_sin = np.sin(out) # (M, D/2) + emb_cos = np.cos(out) # (M, D/2) + + emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) + return emb + + +@get_rope_shape_decorate +@torch.compile(dynamic=True) +def get_rope_shape(org, interpolation_mode, shape): + return ( + F.interpolate( + org.permute((2, 0, 1)).unsqueeze(0), + size=shape, + mode=interpolation_mode, + ) + .squeeze(0) + .permute((1, 2, 0)) + .flatten(end_dim=1) + ) + + +def get_1d_sincos_pos_embed(embed_dim, t_size, cls_token=False): + """ + t_size: int of the temporal size + return: + pos_embed: [t_size, embed_dim] or [1+t_size, embed_dim] (w/ or w/o cls_token) + """ + grid_t = np.arange(t_size, dtype=np.float32) + pos_embed = get_1d_sincos_pos_embed_from_grid(embed_dim, grid_t) + if cls_token: + pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0) + return pos_embed + + +class Learnable2DInterpPosEmbDivided_fixed(nn.Module): + + def __init__( + self, + height: int, + width: int, + num_frames: int, + dim: int, + interpolation_mode: str = "bicubic", + ) -> None: + super().__init__() + self.height = height + self.width = width + self.num_frames = num_frames + self.dim = dim + self.interpolation_mode = interpolation_mode + self.weight = nn.Parameter(torch.empty(height, width, dim)) + self.register_buffer( + "time_weight", + torch.from_numpy(get_1d_sincos_pos_embed(self.dim, self.num_frames)) + .float() + .unsqueeze(1), + persistent=False, + ) + + self.reset_parameters() + + def reset_parameters(self): + nn.init.normal_(self.weight) + + def forward(self, x: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor: + pos_embs = [] + for t, h, w in grid_thws.tolist(): + if t > self.num_frames: + raise ValueError(f"t:{t} > self.num_frames:{self.num_frames}") + if (h, w) == self.weight.shape[:-1]: + pos_emb_2d = self.weight.flatten(end_dim=1) + else: + pos_emb_2d = get_rope_shape( + self.weight, + interpolation_mode=self.interpolation_mode, + shape=(h, w), + ) + + if t == 1: + pos_emb_3d = pos_emb_2d + else: + pos_emb_3d = ( + pos_emb_2d.unsqueeze(0).repeat(t, 1, 1) + self.time_weight[0:t] + ) + + pos_embs.append(pos_emb_3d.reshape(-1, pos_emb_3d.shape[-1])) + + out = x + torch.cat(pos_embs) + return out + + +class Rope2DPosEmbRepeated(nn.Module): + """2D rotary position embedding with multi-resolution support. + + Lifecycle: + 1. At construction, precompute and hold the cis tensor. + 2. Before each forward pass, call ``get_freqs_cis_by_*`` to get the + ``freqs_cis`` tensor for this iteration. + 3. During the forward pass, pass ``freqs_cis`` to each attention layer + and call ``apply`` just before each attention op. Rope is shared + across all attention layers and all heads. + + Refs: + - RoFormer: https://arxiv.org/abs/2104.09864 + - VisionLLaMA: https://arxiv.org/abs/2403.00522 + - https://github.com/Meituan-AutoML/VisionLLaMA/blob/main/dit/models.py + + Args: + dim (int): usually the multi-head attention dimension; must be divisible by 4. + max_height (int): the maximum height of the 2D grid. + max_width (int): the maximum width of the 2D grid. + theta_base (float): the base of the theta. + """ + + def __init__(self, dim: int, max_height: int, max_width: int, theta_base=10000): + super().__init__() + self.dim = dim + if self.dim % 4 != 0: + raise ValueError("dim must be divisible by 4") + self.max_height = max_height + self.max_width = max_width + self.theta_base = theta_base + + def extra_repr(self): + return f"dim={self.dim}, max_height={self.max_height}, max_width={self.max_width}, theta_base={self.theta_base}" + + def _precompute_freqs_cis(self, device: torch.device) -> torch.Tensor: + """Calculate the cis(freqs) for each position in the 2D grid. + Return: complex tensor of shape (max_height, max_width, dim//2) and value: + height axis: ret[h, w, 2*i] = cis(h * theta_base**(-4*i/dim)) + weight axis: ret[h, w, 2*i+1] = cis(w * theta_base**(-4*i/dim)) with (i in [0, dim//4)) + note: `cis` is a mathematical notation defined by cis x = cos x + i sin x, + """ + N = self.max_height * self.max_width + flat_pos = torch.arange(0, N).float().to(device) + x_pos = flat_pos % self.max_width + y_pos = flat_pos // self.max_width + dim_range = ( + torch.arange(0, self.dim, 4)[: (self.dim // 4)].float().to(device) + ) # C/4 + freqs = 1.0 / (self.theta_base ** (dim_range / self.dim)) + x_freqs = torch.outer(x_pos, freqs).float() # N, C/4 + y_freqs = torch.outer(y_pos, freqs).float() # N, C/4 + x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs) # N, C/4 + y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs) # N, C/4 + # N, C/4, 2 + freqs_cis = torch.cat( + [x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1 + ) + # max_height, max_width, C/2 + freqs_cis = freqs_cis.reshape(self.max_height, self.max_width, -1) + return freqs_cis + + def get_freqs_cis( + self, grid_thws: torch.Tensor, device: torch.device + ) -> torch.Tensor: + """ + Args: + grid_thws (torch.Tensor): grid time, height and width + Returns: + freqs_cis: tensor of shape (sum(t * height * width), dim//2) + """ + if not hasattr(self, "freqs_cis"): + self.register_buffer( + "freqs_cis", self._precompute_freqs_cis(device), persistent=False + ) + + shapes = grid_thws.tolist() + if not all( + 1 <= h <= self.max_height and 1 <= w <= self.max_width for t, h, w in shapes + ): + raise ValueError( + f"grid shape out of range: {shapes}, max_height={self.max_height}, " + f"max_width={self.max_width}" + ) + freqs_cis = torch.cat( + [ + self.freqs_cis[:h, :w].reshape(-1, self.dim // 2).repeat(t, 1) + for t, h, w in shapes + ], + dim=0, + ) + return freqs_cis + + +class MoonVision3dPatchEmbed(nn.Module): + + def __init__( + self, + out_dim: int, + in_dim: int = 3, + patch_size: int | tuple[int, int] = (14, 14), + pos_emb_height: int = 14, + pos_emb_width: int = 14, + pos_emb_time: int = 4, + pos_emb_type: str = "divided_fixed", + ): + super().__init__() + if not isinstance(patch_size, int | Sequence): + raise TypeError(f"Invalid patch_size type: {type(patch_size)}") + if isinstance(patch_size, int): + patch_size = (patch_size, patch_size) + if len(patch_size) != 2: + raise ValueError( + f"Expected patch_size to be a tuple of 2, got {patch_size}" + ) + self.patch_size = patch_size + + self.proj = Conv2dLayer( + in_dim, out_dim, kernel_size=patch_size, stride=patch_size + ) + + if pos_emb_type == "divided_fixed": + self.pos_emb = Learnable2DInterpPosEmbDivided_fixed( + height=pos_emb_height, + width=pos_emb_width, + num_frames=pos_emb_time, + dim=out_dim, + ) + else: + raise NotImplementedError(f"Not support pos_emb_type: {pos_emb_type}") + + def forward(self, x: torch.Tensor, grid_thws: torch.Tensor) -> torch.Tensor: + """ + Args: + x (L, Channels): input tensor + grid_thws (N, 3): temporal, height and width + Returns: + (L, Cout) tensor + """ + x = self.proj(x).view(x.size(0), -1) + # apply positional embedding + x = self.pos_emb(x, grid_thws) + return x + + +class MoonViT3dEncoder(nn.Module): + + def __init__( + self, + hidden_dim: int, + num_layers: int, + block_cfg: dict, + video_attn_type: str = "spatial_temporal", + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + + if video_attn_type != "spatial_temporal": + raise ValueError( + f'video_attn_type must be "spatial_temporal", got {video_attn_type}' + ) + self.video_attn_type = video_attn_type + self.rope_2d = Rope2DPosEmbRepeated( + block_cfg["hidden_dim"] // block_cfg["num_heads"], 512, 512 + ) + self.blocks = nn.ModuleList( + [ + MoonViTEncoderLayer( + **block_cfg, + quant_config=quant_config, + prefix=add_prefix(f"blocks.{layer_idx}", prefix), + ) + for layer_idx in range(num_layers) + ] + ) + self.final_layernorm = nn.LayerNorm(hidden_dim) + + def prepare_metadata( + self, grid_thws: torch.Tensor, device: torch.device | None = None + ) -> dict[str, torch.Tensor | int]: + """Eager metadata pass: everything with a GPU->CPU sync or a + data-dependent shape lives here, outside the capture-safe block loop. + + Returns the ``rope_freqs_cis`` / ``cu_seqlens`` tensors plus + ``max_seqlen`` as a Python int (see ``MoonViTEncoderLayer.forward``). + ``max_seqlen`` is materialized numpy-side so the block loop never hits + a ``.item()`` host sync on cudagraph replay. + """ + if device is None: + device = self.final_layernorm.weight.device + rope_freqs_cis = self.rope_2d.get_freqs_cis(grid_thws=grid_thws, device=device) + + grid_thws_np = grid_thws.cpu().numpy() + real_seq_lens = grid_thws_np[:, 0] * grid_thws_np[:, 1] * grid_thws_np[:, 2] + max_seqlen = int(real_seq_lens.max()) if real_seq_lens.size > 0 else 0 + cu_seqlens_np = np.concatenate( + [np.zeros(1, dtype=np.int32), real_seq_lens.cumsum(dtype=np.int32)] + ) + cu_seqlens = torch.from_numpy(cu_seqlens_np).to( + device=device, dtype=torch.int32, non_blocking=True + ) + + return { + "rope_freqs_cis": rope_freqs_cis, + "cu_seqlens": cu_seqlens, + "max_seqlen": max_seqlen, + } + + def forward_blocks( + self, + hidden_states: torch.Tensor, + metadata: dict[str, torch.Tensor | int], + ) -> torch.Tensor: + """Capture-safe encoder body: the block loop + final norm. No host + syncs and no data-dependent control flow, so this region is safe to + record into a CUDA graph. ``metadata`` comes from + :meth:`prepare_metadata`.""" + rope_freqs_cis = metadata["rope_freqs_cis"] + cu_seqlens = metadata["cu_seqlens"] + max_seqlen = metadata["max_seqlen"] + + for block in self.blocks: + hidden_states = block( + hidden_states, cu_seqlens, max_seqlen, rope_freqs_cis=rope_freqs_cis + ) + + return self.final_layernorm(hidden_states) + + +class MoonViT3dPretrainedModel(nn.Module): + def __init__( + self, + config, + mapping: Mapping, + *inputs, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + mm_attention_backend: str | None = None, + **kwargs, + ): + super().__init__() + config = deepcopy(config) + self.config = config + self.merge_kernel_size = config.merge_kernel_size + + self.patch_embed = MoonVision3dPatchEmbed( + out_dim=config.hidden_size, + patch_size=config.patch_size, + pos_emb_height=config.init_pos_emb_height, + pos_emb_width=config.init_pos_emb_width, + pos_emb_time=config.init_pos_emb_time, + pos_emb_type=config.pos_emb_type, + ) + + self.encoder = MoonViT3dEncoder( + hidden_dim=config.hidden_size, + num_layers=config.num_hidden_layers, + block_cfg={ + "num_heads": config.num_attention_heads, + "hidden_dim": config.hidden_size, + "mlp_dim": config.intermediate_size, + "activation": PytorchGELUTanh(), + "attn_bias": True, + "mapping": mapping, + "mm_attention_backend": mm_attention_backend, + }, + video_attn_type=config.video_attn_type, + quant_config=quant_config, + prefix=add_prefix("encoder", prefix), + ) + + @property + def dtype(self) -> torch.dtype: + return self.patch_embed.proj.weight.dtype + + @property + def device(self) -> torch.device: + return self.patch_embed.proj.weight.device + + +class K2VLMultiModalProjector(nn.Module): + """Multi-modal projector with patch merging for K2-VL.""" + + def __init__( + self, + config: KimiK25VisionConfig, + prefix: str = "", + ): + super().__init__() + + # Hidden size after patch merging + merge_h, merge_w = config.merge_kernel_size + self.hidden_size = config.vt_hidden_size * merge_h * merge_w + + self.pre_norm = torch.nn.LayerNorm(config.vt_hidden_size, eps=1e-5) + self.linear_1 = ReplicatedLinear( + self.hidden_size, + self.hidden_size, + bias=True, + prefix=add_prefix("linear_1", prefix), + ) + self.linear_2 = ReplicatedLinear( + self.hidden_size, + config.text_hidden_size, + bias=True, + prefix=add_prefix("linear_2", prefix), + ) + self.act = nn.GELU() + + def forward(self, image_features: torch.Tensor) -> torch.Tensor: + hidden_states = self.pre_norm(image_features).view(-1, self.hidden_size) + hidden_states, _ = self.linear_1(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states, _ = self.linear_2(hidden_states) + return hidden_states + + +@torch.inference_mode() +def mm_projection_auto( + mm_projector: torch.nn.Module | None, vt_output: list[torch.Tensor] +): + """Apply MM projector to vision tower outputs.""" + if mm_projector is None: + return vt_output + + num_embedding_list = [x.shape[0] for x in vt_output] + batched = torch.cat(vt_output, dim=0) + proj_out = mm_projector(batched) + proj_out = proj_out.reshape(-1, proj_out.shape[-1]) + proj_out = torch.split(proj_out, num_embedding_list) + return proj_out + + +class KimiK25ForConditionalGeneration(nn.Module): + def __init__( + self, + config: KimiK25Config, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + is_multimodal_active: bool = True, + mm_attention_backend: str | None = None, + **kwargs, # fix init_tts argument error + ) -> None: + super().__init__() + self.config = config + self.mapping = mapping + self.quant_config = quant_config + self.is_multimodal_active = is_multimodal_active + if not self.is_multimodal_active: + self.vision_tower = None + self.mm_projector = None + else: + self.vision_tower = MoonViT3dPretrainedModel( + config.vision_config, + quant_config=( + quant_config if isinstance(quant_config, ModelSlimConfig) else None + ), + prefix="vision_tower", + mapping=mapping, + mm_attention_backend=mm_attention_backend, + ) + self.mm_projector = K2VLMultiModalProjector(config.vision_config) + + self.language_model = None + if not getattr(config, "encoder_only", False): + self.language_model = DeepseekV3ForCausalLM( + config.text_config, + mapping=mapping, + quant_config=quant_config, + prefix=( + "language_model" + if isinstance(quant_config, (ModelSlimConfig, QuarkConfig)) + else "" + ), + ) + + if self.is_multimodal_active: + # Match vision-tower / mm-projector dtype to language-model dtype; + # the vision tower defaults to float32 while the LM may be bf16 / fp8. + if self.language_model is not None and hasattr( + self.language_model, "dtype" + ): + target_dtype = self.language_model.dtype + self.vision_tower = self.vision_tower.to(dtype=target_dtype) + self.mm_projector = self.mm_projector.to(dtype=target_dtype) + + # image_encoder may be swapped to a cudagraph wrapper by ModelExecutor. + self.vision_embedder = VisionEmbedder() + self.image_encoder = self.get_image_feature + else: + self.vision_embedder = None + self.image_encoder = None + + def get_image_feature(self, items: list[MultimodalDataItem]) -> torch.Tensor: + """Eager image encode via the same ``pre_encode`` / ``forward_blocks`` + / ``post_encode`` decomposition the cudagraph wrapper uses, so the + eager and captured paths share a single source of truth.""" + tokens, grid_thws = self.pre_encode(items) + encoder = self.vision_tower.encoder + encoded = encoder.forward_blocks(tokens, encoder.prepare_metadata(grid_thws)) + # forward_blocks keeps a leading batch dim of 1; squeeze it for + # per-image consumption (mirrors ``out_squeeze_dim=0`` in the + # cudagraph wrapper). + return self.post_encode([encoded.squeeze(0)], grid_thws) + + def pre_encode( + self, items: list[MultimodalDataItem] + ) -> tuple[torch.Tensor, torch.Tensor]: + """Eager patch-embed before the captured region; returns (tokens, grid). + + Reads HF-native ``grid_thws`` on each item (matches the SMG gateway's + Kimi-K2.5 processor). + """ + device = self.vision_tower.device + target_dtype = self.vision_tower.patch_embed.proj.weight.dtype + pixel_values = torch.cat( + [item.feature.to(device, non_blocking=True) for item in items], dim=0 + ).to(dtype=target_dtype) + grid_thws = torch.concat([item.grid_thws for item in items], dim=0).to(device) + hidden_states = self.vision_tower.patch_embed(pixel_values, grid_thws) + return hidden_states, grid_thws + + def post_encode( + self, encoder_outs: list[torch.Tensor], grid_thws: torch.Tensor + ) -> torch.Tensor: + """Eager merge + projection after the captured region; returns features.""" + merged = tpool_patch_merger( + torch.cat(encoder_outs, dim=0), + grid_thws, + merge_kernel_size=self.vision_tower.merge_kernel_size, + ) + proj_out = mm_projection_auto(self.mm_projector, merged) + return torch.cat(proj_out, dim=0) + + def make_encoder_cudagraph_wrappers(self, mapping): + # Captured region is ``MoonViT3dEncoder.forward_blocks`` (token-preserving + # block loop); spatial/temporal merge lives in ``post_encode``, so + # budgets are encoder-input patch counts (``out_div=1``). ``forward_blocks`` + # keeps a leading batch dim of 1 -- ``out_squeeze_dim=0`` drops it before + # per-item slicing. + return { + "image_encoder": EncoderCudaGraphWrapper( + adapter=VisionEncoderCudaGraphAdapter( + tower=self.vision_tower.encoder, + pre_encode=self.pre_encode, + post_encode=self.post_encode, + out_div=1, + merge=1, + input_feature_shape=(self.config.vision_config.hidden_size,), + modality_name="image", + out_squeeze_dim=0, + capture_tp_size=mapping.vision.tp_size, + capture_tp_group=mapping.vision.tp_group, + ), + budget_range=(256, 16384), + ) + } + + def pad_input_ids(self, input_ids: list[int], mm_inputs: MultimodalInputs): + return pad_input_tokens(input_ids, mm_inputs) + + @property + def start_layer(self) -> int: + return self.language_model.start_layer if self.language_model is not None else 0 + + @property + def end_layer(self) -> int: + if self.language_model is not None: + return self.language_model.end_layer + text_config = getattr(self.config, "text_config", None) + return int(getattr(text_config, "num_hidden_layers", 0)) + + @property + def routed_experts_weights_of_layer(self): + return ( + self.language_model._routed_experts_weights_of_layer.value + if self.language_model is not None + else {} + ) + + @torch.no_grad() + def multimodal_input_embeds( + self, + input_ids: torch.Tensor, + ctx, + multimodal_context, + ) -> torch.Tensor | None: + """Merged text+vision input embeddings, or ``None`` for a plain text step. + + Kimi-K2.5's multimodal path is embeds-only -- the vision features are + scattered into the input embeddings and nothing else reaches the + language model (no per-layer extras like deepstack) -- so both the + eager ``forward`` below and a prefill-graph replay take the exact same + tensor from here. + """ + if ( + multimodal_context is None + or self.vision_embedder is None + or not multimodal_context.has_extend_inputs() + or ctx.forward_mode.is_decode_or_idle() + ): + return None + input_embeds, model_kwargs = self.vision_embedder.apply( + input_ids=input_ids, + text_embedding=self.get_input_embeddings(), + ctx=multimodal_context, + encoders={Modality.IMAGE: EncoderSpec(self.image_encoder)}, + multimodal_model=self, + is_decode_or_idle=ctx.forward_mode.is_decode_or_idle(), + ) + assert not model_kwargs, "Kimi-K2.5 multimodal path must stay embeds-only" + return input_embeds + + def forward( + self, + ctx, + input_ids: torch.Tensor, + positions: torch.Tensor, + out_cache_loc: torch.Tensor, + **kwargs, + ): + if self.language_model is None: + raise RuntimeError("KimiK25 language_model is not initialized.") + multimodal_context = kwargs.pop("multimodal_context", None) + input_embeds = self.multimodal_input_embeds(input_ids, ctx, multimodal_context) + if input_embeds is not None: + kwargs["input_embeds"] = input_embeds + return self.language_model.forward( + ctx, + input_ids, + positions, + out_cache_loc, + **kwargs, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + """Load weights for the model, separating vision and language weights""" + vision_weights = [] + language_weights = [] + + for name, loaded_weight in weights: + # nvidia/Kimi-K2.5-NVFP4 stores decoder layers under + # language_model.layers.*, while TokenSpeed's DeepSeek module + # expects model.layers.* after stripping language_model. + if name.startswith("language_model.layers."): + name = name.replace( + "language_model.layers.", "language_model.model.layers.", 1 + ) + + if "vision_tower" in name or "mm_projector" in name: + name = name.replace(r"wqkv.", r"attn.qkv_proj.") + name = name.replace(r"wo.", r"attn.proj.") + name = name.replace("mm_projector.proj.0", "mm_projector.linear_1") + name = name.replace("mm_projector.proj.2", "mm_projector.linear_2") + vision_weights.append((name, loaded_weight)) + else: + name = name.replace("language_model.", "") + language_weights.append((name, loaded_weight)) + + if self.is_multimodal_active and not getattr( + self.config, "language_only", False + ): + vision_state_dict = dict(vision_weights) + params_dict = dict(self.named_parameters(remove_duplicate=False)) + for name, loaded_weight in vision_state_dict.items(): + if name not in params_dict: + raise ValueError(f"Weight {name} not found in params_dict") + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + if not getattr(self.config, "encoder_only", False) and language_weights: + self.language_model.load_weights(language_weights) + + @classmethod + def get_model_config_for_expert_location(cls, config: KimiK25Config): + text_config = config.text_config + return ModelConfigForExpertLocation( + num_layers=text_config.num_hidden_layers, + num_logical_experts=text_config.n_routed_experts, + num_groups=text_config.n_group, + ) + + def set_eagle3_layers_to_capture(self, layer_ids: list[int] | None = None) -> None: + """Set the layers to capture for EAGLE3 speculative decoding.""" + if self.language_model is None or not hasattr( + self.language_model, "set_eagle3_layers_to_capture" + ): + raise AttributeError( + "language_model does not support EAGLE3 speculative decoding." + ) + + self.language_model.set_eagle3_layers_to_capture(layer_ids) + + def set_dflash_layers_to_capture(self, layer_ids: list[int]) -> None: + """Set the layers to capture for DFLASH draft model training.""" + if not hasattr(self.language_model, "set_dflash_layers_to_capture"): + raise AttributeError( + "language_model does not support DFLASH layer capture." + ) + + self.language_model.set_dflash_layers_to_capture(layer_ids) + + def get_input_embeddings(self): + if hasattr(self.language_model, "get_input_embeddings"): + return self.language_model.get_input_embeddings() + if hasattr(self.language_model, "model") and hasattr( + self.language_model.model, "embed_tokens" + ): + return self.language_model.model.embed_tokens + raise AttributeError("language_model does not support get_input_embeddings().") + + @property + def lm_head(self): + if not hasattr(self.language_model, "lm_head"): + raise AttributeError("language_model does not expose lm_head.") + + return self.language_model.lm_head + + @property + def logits_processor(self): + if self.language_model is None or not hasattr( + self.language_model, "logits_processor" + ): + raise AttributeError("language_model does not expose logits_processor.") + + return self.language_model.logits_processor + + def get_embed_and_head(self) -> tuple[torch.Tensor, torch.Tensor]: + """Get embedding and LM head weights for speculative decoding.""" + if self.language_model is None or not hasattr( + self.language_model, "get_embed_and_head" + ): + raise AttributeError( + "language_model does not support get_embed_and_head()." + ) + + return self.language_model.get_embed_and_head() + + def set_embed_and_head(self, embed: torch.Tensor, head: torch.Tensor) -> None: + """Set embedding and LM head weights for speculative decoding.""" + if self.language_model is None or not hasattr( + self.language_model, "set_embed_and_head" + ): + raise AttributeError( + "language_model does not support set_embed_and_head()." + ) + + self.language_model.set_embed_and_head(embed, head) + + +EntryClass = [KimiK25ForConditionalGeneration] diff --git a/python/tokenspeed/runtime/models/llama.py b/python/tokenspeed/runtime/models/llama.py new file mode 100644 index 0000000..5d98f5d --- /dev/null +++ b/python/tokenspeed/runtime/models/llama.py @@ -0,0 +1,395 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Inference-only dense Llama model compatible with HuggingFace weights. + +Covers Llama-2 / Llama-3 / Llama-3.1 / Llama-3.2 dense checkpoints whose +``config.architectures`` is ``["LlamaForCausalLM"]``. MoE and Eagle3 draft +variants have their own modules (``longcat_large.py``, ``llama_eagle3.py``). +""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +import torch +from torch import nn +from transformers import LlamaConfig + +from tokenspeed.runtime.configs.utils import get_rope_theta +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.layers.activation import SiluAndMul +from tokenspeed.runtime.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from tokenspeed.runtime.layers.paged_attention import PagedAttention +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.layers.rotary_embedding import get_rope +from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader +from tokenspeed.runtime.models.base import ( + BaseCausalLM, + BaseDecoderLayer, + BaseTransformerModel, +) +from tokenspeed.runtime.models.utils import ( + create_fused_set_kv_buffer_arg, + validate_attention_partition, +) +from tokenspeed.runtime.utils import add_prefix +from tokenspeed.runtime.utils.pdl import pdl_enabled + + +class LlamaMLP(nn.Module): + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + tp_rank = mapping.dense.tp_rank + tp_size = mapping.dense.tp_size + tp_group = mapping.dense.tp_group + + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + tp_rank=tp_rank, + tp_size=tp_size, + tp_group=tp_group, + prefix=add_prefix("gate_up_proj", prefix), + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=False, + tp_rank=tp_rank, + tp_size=tp_size, + tp_group=tp_group, + prefix=add_prefix("down_proj", prefix), + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. Only silu is supported." + ) + self.act_fn = SiluAndMul() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if x.shape[0] == 0: + return x + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class LlamaAttention(nn.Module): + + def __init__( + self, + config: LlamaConfig, + mapping: Mapping, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + layer_id: int = 0, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + qkv_input_size: int | None = None, + ) -> None: + super().__init__() + self.hidden_size = hidden_size + + self.attn_tp_size = mapping.attn.tp_size + self.attn_tp_rank = mapping.attn.tp_rank + attn_tp_group = mapping.attn.tp_group + + self.total_num_heads = num_heads + self.total_num_kv_heads = num_kv_heads + validate_attention_partition( + self.total_num_heads, + self.total_num_kv_heads, + self.attn_tp_size, + ) + self.num_heads = self.total_num_heads // self.attn_tp_size + self.num_kv_heads = max(1, self.total_num_kv_heads // self.attn_tp_size) + self.head_dim = getattr( + config, "head_dim", self.hidden_size // self.total_num_heads + ) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + + rope_theta = get_rope_theta(config) + rope_scaling = getattr(config, "rope_scaling", None) + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + # Dense Llama is consistently bias-free (`attention_bias=False` in every + # upstream release). Still read it off the config so forks that flip + # the flag load without surprises. + attention_bias = getattr(config, "attention_bias", False) + + self.qkv_proj = QKVParallelLinear( + qkv_input_size or hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=attention_bias, + quant_config=quant_config, + tp_rank=self.attn_tp_rank, + tp_size=self.attn_tp_size, + tp_group=attn_tp_group, + prefix=add_prefix("qkv_proj", prefix), + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=attention_bias, + quant_config=quant_config, + reduce_results=False, + tp_rank=self.attn_tp_rank, + tp_size=self.attn_tp_size, + tp_group=attn_tp_group, + prefix=add_prefix("o_proj", prefix), + ) + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + ) + self.attn = PagedAttention( + self.num_heads, + self.head_dim, + self.head_dim**-0.5, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + ) -> torch.Tensor: + # Skip the QKV projection, RoPE, attention, and o_proj kernels when + # the batch row is empty (e.g. idle ranks under DP attention). Matches + # the short-circuit ``LlamaMLP.forward`` already has. + if hidden_states.shape[0] == 0: + return hidden_states.new_zeros( + (0, self.hidden_size), dtype=hidden_states.dtype + ) + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + attn_output = self._attn(positions, q, k, v, ctx, out_cache_loc) + output, _ = self.o_proj(attn_output) + return output + + def _attn( + self, + positions: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + ) -> torch.Tensor: + """RoPE + attention (pre-o_proj), with optional fused KV pre-write. + + When the backend supports KV pre-write *and* ``create_fused_set_kv_buffer_arg`` + accepts the layer's scales, fused rope writes KV directly into the cache + so the attention call can run with ``save_kv_cache=False`` (saves one + kernel launch). Otherwise we fall back to plain RoPE + ``self.attn(q, k, v)`` + so the backend writes KV the normal way — without this fallback, layers + with non-trivial k/v scales silently lose their KV writes. Subclasses + (e.g. Eagle3 draft head) override this hook to insert spec-decode + behaviour around the same scaffolding. + """ + if ctx.attn_backend.support_kv_cache_prewrite(ctx.forward_mode): + fused_kv_arg = self._build_fused_kv_arg(v, ctx, out_cache_loc) + if fused_kv_arg is not None: + q_rope = self._fused_rope_kv_write(positions, q, k, fused_kv_arg) + return self.attn( + q_rope, + None, + None, + save_kv_cache=False, + ctx=ctx, + out_cache_loc=out_cache_loc, + ) + q, k = self.rotary_emb(positions, q, k) + return self.attn(q, k, v, ctx=ctx, out_cache_loc=out_cache_loc) + + def _build_fused_kv_arg( + self, + v: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + ): + """Try to build the fused RoPE+KV-write descriptor; returns ``None`` if + the helper rejects the layer (e.g. non-trivial k/v scales).""" + n = v.shape[0] + return create_fused_set_kv_buffer_arg( + value=v.view(n, self.num_kv_heads, self.head_dim), + layer=self.attn, + out_cache_loc=out_cache_loc, + token_to_kv_pool=ctx.token_to_kv_pool, + ) + + def _fused_rope_kv_write( + self, + positions: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + fused_kv_arg, + ) -> torch.Tensor: + """Fused RoPE that writes KV into cache (via ``fused_kv_arg``) and + returns the rope'd Q.""" + n = q.shape[0] + q_rope = torch.empty((n, self.q_size), dtype=q.dtype, device=q.device) + self.rotary_emb( + positions, + q, + k, + fused_set_kv_buffer_arg=fused_kv_arg, + output_q_rope=q_rope, + enable_pdl=pdl_enabled(), + ) + return q_rope + + +class LlamaDecoderLayer(BaseDecoderLayer): + + def resolve_attn(self, prefix: str) -> nn.Module: + return LlamaAttention( + config=self.config, + mapping=self.mapping, + hidden_size=self.config.hidden_size, + num_heads=self.config.num_attention_heads, + num_kv_heads=self.config.num_key_value_heads, + layer_id=self.layer_id, + quant_config=self.quant_config, + prefix=add_prefix("self_attn", prefix), + ) + + def resolve_mlp(self, prefix: str) -> nn.Module: + return LlamaMLP( + hidden_size=self.config.hidden_size, + intermediate_size=self.config.intermediate_size, + hidden_act=self.config.hidden_act, + mapping=self.mapping, + quant_config=self.quant_config, + prefix=add_prefix("mlp", prefix), + ) + + +class LlamaModel(BaseTransformerModel): + layer_cls = LlamaDecoderLayer + + +class LlamaForCausalLM(BaseCausalLM): + model_cls = LlamaModel + + # BitsAndBytes target/stacked modules — kept in sync with the Qwen3 / MoE + # variants so a single quantization config works across the Llama family. + default_bitsandbytes_target_modules = [ + ".gate_proj.", + ".down_proj.", + ".up_proj.", + ".q_proj.", + ".k_proj.", + ".v_proj.", + ".o_proj.", + ] + bitsandbytes_stacked_params_mapping = { + "q_proj": ("qkv_proj", 0), + "k_proj": ("qkv_proj", 1), + "v_proj": ("qkv_proj", 2), + "gate_proj": ("gate_up_proj", 0), + "up_proj": ("gate_up_proj", 1), + } + + def get_stacked_params_mapping(self) -> list[tuple[str, str, int | str]]: + return [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + def load_weights( + self, weights: Iterable[tuple[str, torch.Tensor]], **kwargs: Any + ) -> None: + stacked_params_mapping = self.get_stacked_params_mapping() + params_dict = dict(self.named_parameters()) + tie_word_embeddings = getattr(self.config, "tie_word_embeddings", False) + + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + continue + # Llama-3.2-1B / 3B ship with tied input+output embeddings — some HF + # checkpoint variants still serialize lm_head.weight, skip it so we + # don't double-load into the shared embed_tokens parameter. + if tie_word_embeddings and "lm_head.weight" in name: + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + if name.endswith(".bias") and name not in params_dict: + continue + if name not in params_dict: + continue + param = params_dict[name] + # Fused q/k/v and gate/up parameters are built by distributed + # linear layers that install ``weight_loader`` during init; the + # ``getattr`` fallback just guards against stray non-fused + # parameters that happened to match the pattern (e.g. a user + # fork that registers a plain ``qkv_proj`` buffer). + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight, shard_id) + break + else: + if name.endswith(".bias") and name not in params_dict: + continue + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + +EntryClass = LlamaForCausalLM diff --git a/python/tokenspeed/runtime/models/llama_eagle3.py b/python/tokenspeed/runtime/models/llama_eagle3.py new file mode 100644 index 0000000..11140c2 --- /dev/null +++ b/python/tokenspeed/runtime/models/llama_eagle3.py @@ -0,0 +1,660 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""LLaMA Eagle3 draft model for speculative decoding. + +Extends base classes. Preserves the low-latency fused allreduce+norm +path from the original implementation. +""" + +from __future__ import annotations + +from collections.abc import Iterable + +import torch +from torch import nn +from transformers import LlamaConfig + +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.execution.context import ( + ForwardContext, + report_collective_sizing, +) +from tokenspeed.runtime.execution.forward_batch_info import ForwardMode +from tokenspeed.runtime.layers.activation import SiluAndMul +from tokenspeed.runtime.layers.common import concat +from tokenspeed.runtime.layers.layernorm import RMSNorm +from tokenspeed.runtime.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + RowParallelLinear, +) +from tokenspeed.runtime.layers.logits_processor import LogitsProcessor +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.layers.vocab_parallel_embedding import ParallelLMHead +from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader +from tokenspeed.runtime.models.base import ( + BaseCausalLM, + BaseDecoderLayer, + BaseTransformerModel, +) +from tokenspeed.runtime.models.llama import LlamaAttention as BaseLlamaAttention +from tokenspeed.runtime.utils import add_prefix, get_colorful_logger + +logger = get_colorful_logger(__name__) + + +# --------------------------------------------------------------------------- +# Attention +# --------------------------------------------------------------------------- + + +class LlamaAttention(BaseLlamaAttention): + """Eagle3 draft head attention. + + Inherits ``__init__`` (with ``qkv_input_size=2*hidden_size`` for the + [embed || hidden] concat) and ``forward`` (= qkv_proj + o_proj scaffolding) + from base. Overrides ``_attn`` so the draft's first step skips dead + catch-up rows: on backends that support fused KV pre-write, q is sliced + to one live row per request and dispatched as DECODE; otherwise the + fallback runs the full N-row attn and post-slices the output. Inactive + draft steps delegate to base. + """ + + def _attn( + self, + positions: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + ) -> torch.Tensor: + # Active draft first step (drafter set up gather_ids + accept_lengths). + # Covers both decode catch-up and prefill catch-up; multi-step decode + # delegates to base. + if ctx.accept_lengths is None: + return super()._attn(positions, q, k, v, ctx, out_cache_loc) + + if ctx.attn_backend.support_kv_cache_prewrite(ctx.forward_mode): + fused_kv_arg = self._build_fused_kv_arg(v, ctx, out_cache_loc) + if fused_kv_arg is not None: + # Trim only on the sliced single-token decode path; the + # post-slice fallback below still runs full N-row attn and + # needs the original seq_lens. + self._apply_correction(ctx) + q_rope = self._fused_rope_kv_write( + positions, q, k, fused_kv_arg + ).index_select(0, ctx.gather_ids) + # record_kv_cache (keyed off the real mode) forces the backend's + # PD layerwise cache-step record that the DECODE dispatch would + # otherwise skip on an EXTEND/MIXED catch-up. + return ctx.attn_backend.forward( + q_rope, + None, + None, + self.attn, + out_cache_loc, + ctx.token_to_kv_pool, + ForwardMode.DECODE, + ctx.bs, + save_kv_cache=False, + record_kv_cache=not ctx.forward_mode.is_decode_or_idle(), + ) + q, k = self.rotary_emb(positions, q, k) + return self.attn(q, k, v, ctx=ctx, out_cache_loc=out_cache_loc).index_select( + 0, ctx.gather_ids + ) + + def _apply_correction(self, ctx: ForwardContext) -> None: + """Trim decode rows' cache_seqlens by ``spec_num_tokens - accept_lengths``.""" + seq_lens_buf = ctx.draft_seq_lens_buf + if seq_lens_buf is None or ctx.accept_lengths is None: + return + num_extends = ctx.num_extends + if num_extends >= ctx.bs: + return + correction = ( + ctx.attn_backend.spec_num_tokens - ctx.accept_lengths[num_extends:] + ).to(seq_lens_buf.dtype) + seq_lens_buf[num_extends : ctx.bs].sub_(correction).clamp_(min=1) + + +# --------------------------------------------------------------------------- +# MLP +# --------------------------------------------------------------------------- + + +class LlamaMLP(nn.Module): + + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + + super().__init__() + tp_rank = mapping.dense.tp_rank + tp_size = mapping.dense.tp_size + tp_group = mapping.dense.tp_group + + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + tp_size=tp_size, + tp_rank=tp_rank, + tp_group=tp_group, + prefix=add_prefix("gate_up_proj", prefix), + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=False, + tp_rank=tp_rank, + tp_size=tp_size, + tp_group=tp_group, + prefix=add_prefix("down_proj", prefix), + ) + self.act_fn = SiluAndMul() + self.gateup_unquanted = quant_config is None + + def forward(self, x, block_scale=None): + + if x.shape[0] == 0: + return x + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +# --------------------------------------------------------------------------- +# Decoder layer +# --------------------------------------------------------------------------- + + +class Eagle3DecoderLayer(BaseDecoderLayer): + """Eagle3 decoder layer with low-latency fused allreduce+norm path. + + Inherits norm/attn/mlp/comm_manager init from BaseDecoderLayer. + Overrides forward with eagle3-specific embed+hidden concat logic. + """ + + def __init__( + self, + config: LlamaConfig, + layer_id: int, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + + self._eagle3_config = config + self._eagle3_mapping = mapping + self._eagle3_quant_config = quant_config + self._eagle3_prefix = prefix + + super().__init__( + config=config, + layer_id=layer_id, + mapping=mapping, + quant_config=quant_config, + prefix=prefix, + ) + + self.hidden_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def resolve_attn(self, prefix: str) -> nn.Module: + + config = self._eagle3_config + return LlamaAttention( + config, + self._eagle3_mapping, + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + layer_id=self.layer_id, + quant_config=self._eagle3_quant_config, + prefix=add_prefix("self_attn", prefix), + qkv_input_size=2 * config.hidden_size, + ) + + def resolve_mlp(self, prefix: str) -> nn.Module: + + config = self._eagle3_config + inter_size = ( + config.intermediate_size_mlp + if config.model_type == "llama4_text" + else config.intermediate_size + ) + + return LlamaMLP( + config.hidden_size, + inter_size, + config.hidden_act, + self._eagle3_mapping, + self._eagle3_quant_config, + prefix=f"{prefix}.mlp", + ) + + def _maybe_narrow_residual( + self, + residual: torch.Tensor, + ctx: ForwardContext, + ) -> torch.Tensor: + """Align residual with attn output narrowed to [bs, H].""" + if ctx.accept_lengths is not None and not ctx.forward_mode.is_idle(): + return residual.index_select(0, ctx.gather_ids) + return residual + + def forward_low_latency( + self, + positions: torch.Tensor, + embeds: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + residual: torch.Tensor | None, + final_norm: RMSNorm = None, + fuse_embed_reduce: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + residual = hidden_states + + if fuse_embed_reduce: + # Fuse embedding allreduce with input_layernorm. + embeds, _, *_ = self.input_layernorm.forward_with_allreduce_fusion( + self.mapping.attn.tp_rank, + self.mapping.attn.tp_group, + embeds, + torch.zeros_like(embeds), + ) + else: + embeds = self.input_layernorm(embeds) + + hidden_states = self.hidden_norm(hidden_states) + hidden_states = concat(embeds, hidden_states) + + # Attention + hidden_states = self.comm_manager.pre_attn_comm(hidden_states, ctx) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ctx=ctx, + out_cache_loc=out_cache_loc, + ) + residual = self._maybe_narrow_residual(residual, ctx) + + # Fused post-attn allreduce + norm (uses attn tp group) + block_scale = None + hidden_states, residual, block_scale, *_ = ( + self.post_attention_layernorm.forward_with_allreduce_fusion( + self.mapping.attn.tp_rank, + self.mapping.attn.tp_group, + hidden_states, + residual, + fuse_block_quant_fp8=not self.mlp.gateup_unquanted, + ) + ) + + hidden_states = self.mlp(hidden_states, block_scale) + + # Fused final allreduce + norm (uses dense tp group) + hidden_states, residual, *_ = final_norm.forward_with_allreduce_fusion( + self.mapping.dense.tp_rank, + self.mapping.dense.tp_group, + hidden_states, + residual, + fuse_block_quant_fp8=False, + ) + + return hidden_states, residual + + def forward( + self, + positions: torch.Tensor, + embeds: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + residual: torch.Tensor | None, + final_norm: RMSNorm = None, + fuse_embed_reduce: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: + + if self.comm_manager.should_fuse(hidden_states.shape[0]): + return self.forward_low_latency( + positions, + embeds, + hidden_states, + ctx, + out_cache_loc, + residual, + final_norm, + fuse_embed_reduce=fuse_embed_reduce, + ) + + # Non-fused path: fuse_embed_reduce is always False here because + # the model only sets it when should_fuse() is True. + residual = hidden_states + embeds = self.input_layernorm(embeds) + hidden_states = self.hidden_norm(hidden_states) + hidden_states = torch.cat([embeds, hidden_states], dim=-1) + + # Attention + hidden_states = self.comm_manager.pre_attn_comm(hidden_states, ctx) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ctx=ctx, + out_cache_loc=out_cache_loc, + ) + residual = self._maybe_narrow_residual(residual, ctx) + hidden_states, residual = self.comm_manager.post_attn_comm( + hidden_states, residual, ctx + ) + + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + + # MLP + hidden_states = self.comm_manager.pre_mlp_comm(hidden_states, ctx) + hidden_states = self.mlp(hidden_states) + hidden_states, residual = self.comm_manager.post_mlp_comm( + hidden_states, residual, ctx + ) + + return hidden_states, residual + + +# --------------------------------------------------------------------------- +# Model and CausalLM +# --------------------------------------------------------------------------- + + +class Eagle3LlamaModel(BaseTransformerModel): + + layer_cls = Eagle3DecoderLayer + + def __init__( + self, + config: LlamaConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + + super().__init__( + config=config, + mapping=mapping, + quant_config=quant_config, + prefix=prefix, + ) + + # Eagle3 uses "midlayer" (not "layers.0") in checkpoint weights. + # Re-register the single layer under the correct name. + self.midlayer = self.layers[0] + del self.layers + + self.num_fc_input_dim = ( + len(config.eagle_aux_hidden_state_layer_ids) + if hasattr(config, "eagle_aux_hidden_state_layer_ids") + else 3 + ) + + self.fc = ColumnParallelLinear( + config.hidden_size * self.num_fc_input_dim, + config.hidden_size, + bias=False, + gather_output=True, + quant_config=quant_config, + prefix=add_prefix("fc", prefix), + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + + # norm_before_fc: RMSNorm over the concatenated aux states before fc (replicated) + self.input_norm = ( + RMSNorm( + config.hidden_size * self.num_fc_input_dim, + eps=config.rms_norm_eps, + ) + if getattr(config, "norm_before_fc", False) + else None + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + input_embeds: torch.Tensor = None, + hidden_states: torch.Tensor = None, + ) -> torch.Tensor: + + if input_embeds is None: + # When TP > 1 and fused allreduce+norm is available, skip the + # NCCL allreduce in the embedding and let the midlayer fuse it + # with the input_layernorm. + midlayer = self.midlayer + num_tokens = input_ids.shape[0] + fuse_embed_reduce = ( + self.mapping.attn.tp_size > 1 + and midlayer.comm_manager.should_fuse(num_tokens) + ) + embeds = self.embed_tokens(input_ids, reduce_results=not fuse_embed_reduce) + else: + embeds = input_embeds + fuse_embed_reduce = False + + if hidden_states is None: + raise ValueError("Eagle3 forward requires hidden_states") + + if hidden_states.size(-1) != embeds.size(-1): + if self.input_norm is not None: + hidden_states = self.input_norm(hidden_states) + hidden_states, _ = self.fc(hidden_states) + + residual = None + midlayer = self.midlayer + hidden_states, residual = midlayer( + positions, + embeds, + hidden_states, + ctx, + out_cache_loc, + residual, + self.norm, + fuse_embed_reduce=fuse_embed_reduce, + ) + + # Decide on pre-slice token count so this matches the path midlayer + # actually took; under draft reduce, hidden_states.shape[0] shrinks. + if midlayer.comm_manager.should_fuse(input_ids.shape[0]): + hidden_states_to_logits, hidden_states_to_aux = hidden_states, residual + else: + hidden_states_to_logits, hidden_states_to_aux = self.norm( + hidden_states, residual + ) + hidden_states_to_logits, _ = midlayer.comm_manager.post_final_norm_comm( + hidden_states_to_logits, None, ctx + ) + hidden_states_to_aux, _ = midlayer.comm_manager.post_final_norm_comm( + hidden_states_to_aux, None, ctx + ) + + return hidden_states_to_logits, [hidden_states_to_aux] + + +class LlamaForCausalLMEagle3(BaseCausalLM): + + model_cls = Eagle3LlamaModel + + def __init__( + self, + config: LlamaConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + + nn.Module.__init__(self) + self.config = config + self.mapping = mapping + self.quant_config = quant_config + + if self.config.num_hidden_layers != 1: + raise ValueError("EAGLE3 currently only supports 1 layer") + + self.model = self.resolve_model(config, mapping, quant_config, prefix) + + self.load_lm_head_from_target = False + if self.config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + if getattr(config, "draft_vocab_size", None) is None: + self.load_lm_head_from_target = True + self.lm_head = ParallelLMHead( + getattr(config, "draft_vocab_size", None) or config.vocab_size, + config.hidden_size, + quant_config=quant_config, + tp_rank=mapping.attn.tp_rank, + tp_size=mapping.attn.tp_size, + tp_group=mapping.attn.tp_group, + prefix=add_prefix("lm_head", prefix), + ) + + self.logits_processor = LogitsProcessor( + config, + skip_all_gather=self.mapping.attn.has_dp, + do_argmax=True, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + self.capture_aux_hidden_states = True + self.hot_token_id = None + + def forward( + self, + ctx: ForwardContext, + input_ids: torch.Tensor, + positions: torch.Tensor, + out_cache_loc: torch.Tensor, + **kwargs, + ) -> torch.Tensor: + with report_collective_sizing(ctx, ctx.bs, ctx.global_bs): + return super().forward(ctx, input_ids, positions, out_cache_loc, **kwargs) + + def prepare_model_kwargs( + self, ctx: ForwardContext, input_ids: torch.Tensor, kwargs: dict + ) -> dict: + model_kwargs = super().prepare_model_kwargs(ctx, input_ids, kwargs) + captured_hidden_states = kwargs.get("captured_hidden_states") + if captured_hidden_states is not None: + model_kwargs["hidden_states"] = captured_hidden_states + else: + # During CUDA graph capture warmup, provide dummy hidden states. + num_tokens = input_ids.shape[0] + hidden_size = self.config.hidden_size + num_fc = self.model.num_fc_input_dim + model_kwargs["hidden_states"] = torch.zeros( + num_tokens, + hidden_size * num_fc, + dtype=torch.bfloat16, + device=input_ids.device, + ) + return model_kwargs + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> None: + + params_dict = dict(self.named_parameters()) + stacked_params_mapping = [ + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + + for name, loaded_weight in weights: + # some Eagle3 checkpoints name the block "layers.0" not "midlayer" + if name.startswith("layers.0."): + name = "midlayer." + name[len("layers.0.") :] + + if "d2t" in name: + self.hot_token_id = loaded_weight + torch.arange(loaded_weight.shape[0]) + continue + + if "t2d" in name: + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + param_name = f"model.{name}" if name not in params_dict else name + if param_name in params_dict: + param = params_dict[param_name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight, shard_id) + break + else: + param_name = name if name in params_dict else f"model.{name}" + if param_name in params_dict: + param = params_dict[param_name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + + def get_hot_token_id(self): + return self.hot_token_id + + def get_embed(self): + return self.model.embed_tokens.weight + + def set_embed_and_head(self, embed, head): + # If draft hidden size != target hidden size, embed cannot be shared + if ( + hasattr(self.config, "target_hidden_size") + and self.config.target_hidden_size != self.config.hidden_size + ): + return + del self.model.embed_tokens.weight + self.model.embed_tokens.weight = embed + if head is not None and self.load_lm_head_from_target: + del self.lm_head.weight + self.lm_head.weight = head + torch.cuda.empty_cache() + torch.cuda.synchronize() + + +EntryClass = [LlamaForCausalLMEagle3] diff --git a/python/tokenspeed/runtime/models/longcat_flash.py b/python/tokenspeed/runtime/models/longcat_flash.py new file mode 100644 index 0000000..214ee65 --- /dev/null +++ b/python/tokenspeed/runtime/models/longcat_flash.py @@ -0,0 +1,913 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from collections.abc import Iterable as _Iterable + +import torch +import torch.nn as nn +import torch.nn.functional as _F +from tokenspeed_kernel.platform import current_platform as _current_platform +from tokenspeed_kernel.thirdparty.cuda import dsv3_router_gemm as _dsv3_router_gemm +from tokenspeed_kernel.thirdparty.cuda import ( + moe_finalize_fuse_shared as _moe_finalize_fuse_shared, +) +from transformers import PretrainedConfig as _PretrainedConfig + +from tokenspeed.runtime.configs.utils import get_rope_theta as _get_rope_theta +from tokenspeed.runtime.distributed.comm_manager import CommManager as _CommManager +from tokenspeed.runtime.distributed.mapping import Mapping as _Mapping +from tokenspeed.runtime.execution.context import ForwardContext as _ForwardContext +from tokenspeed.runtime.execution.cuda_graph_wrapper import ( + get_is_capture_mode as _get_is_capture_mode, +) +from tokenspeed.runtime.layers.layernorm import RMSNorm as _RMSNorm +from tokenspeed.runtime.layers.linear import ReplicatedLinear +from tokenspeed.runtime.layers.moe import ( + ExpertCheckpointSchema as _ExpertCheckpointSchema, +) +from tokenspeed.runtime.layers.moe import ( + build_moe_checkpoint_loader as _build_moe_checkpoint_loader, +) +from tokenspeed.runtime.layers.moe.expert import MoELayer as _MoELayer +from tokenspeed.runtime.layers.moe.topk import TopK as _TopK +from tokenspeed.runtime.layers.moe.topk import TopKOutputFormat as _TopKOutputFormat +from tokenspeed.runtime.layers.moe.utils import RoutingMethodType as _RoutingMethodType +from tokenspeed.runtime.layers.quantization.base_config import ( + QuantizationConfig as _QuantizationConfig, +) +from tokenspeed.runtime.layers.quantization.utils import block_dequant as _block_dequant +from tokenspeed.runtime.layers.quantization.utils import ( + should_ignore_quant_layer as _should_ignore_quant_layer, +) +from tokenspeed.runtime.layers.utils import get_layer_id as _get_layer_id +from tokenspeed.runtime.layers.vocab_parallel_embedding import ( + VocabParallelEmbedding as _VocabParallelEmbedding, +) +from tokenspeed.runtime.model_loader.weight_utils import ( + default_weight_loader as _default_weight_loader, +) +from tokenspeed.runtime.model_loader.weight_utils import ( + kv_cache_scales_loader as _kv_cache_scales_loader, +) +from tokenspeed.runtime.models.base import BaseCausalLM as _BaseCausalLM +from tokenspeed.runtime.models.deepseek_v3 import ( + DeepseekV3AttentionMLA as _DeepseekV3AttentionMLA, +) +from tokenspeed.runtime.models.deepseek_v3 import DeepseekV3MLP as _DeepseekV3MLP +from tokenspeed.runtime.moe.distribution_recorder import ( + get_global_expert_distribution_recorder as _get_global_expert_distribution_recorder, +) +from tokenspeed.runtime.moe.expert_location import ( + ModelConfigForExpertLocation as _ModelConfigForExpertLocation, +) +from tokenspeed.runtime.utils import LazyValue, add_prefix, get_colorful_logger +from tokenspeed.runtime.utils.cuda_stream import StreamFork as _StreamFork +from tokenspeed.runtime.utils.env import global_server_args_dict +from tokenspeed.runtime.utils.pdl import pdl_enabled as _pdl_enabled + +_longcat_logger = get_colorful_logger(__name__) +_longcat_platform = _current_platform() +_longcat_is_hopper_plus = _longcat_platform.is_hopper_plus +_LONGCAT_OPTIONAL_MISSING_WEIGHT_SUFFIXES = ( + ".k_scale", + ".v_scale", +) + + +def _ensure_longcat_config(config): + """Normalize LongCat HF config aliases used by the runtime layers.""" + + if not hasattr(config, "num_hidden_layers") and hasattr(config, "num_layers"): + config.num_hidden_layers = config.num_layers + if not hasattr(config, "intermediate_size") and hasattr(config, "ffn_hidden_size"): + config.intermediate_size = config.ffn_hidden_size + if not hasattr(config, "moe_intermediate_size"): + if hasattr(config, "expert_ffn_hidden_size"): + config.moe_intermediate_size = config.expert_ffn_hidden_size + else: + config.moe_intermediate_size = config.intermediate_size + if not hasattr(config, "num_experts_per_tok") and hasattr(config, "moe_topk"): + config.num_experts_per_tok = config.moe_topk + if not hasattr(config, "moe_topk") and hasattr(config, "num_experts_per_tok"): + config.moe_topk = config.num_experts_per_tok + + if not hasattr(config, "hidden_act"): + config.hidden_act = "silu" + if not hasattr(config, "norm_topk_prob"): + config.norm_topk_prob = False + if not hasattr(config, "zero_expert_num"): + config.zero_expert_num = 0 + if not hasattr(config, "zero_expert_type"): + config.zero_expert_type = "" + if not hasattr(config, "router_bias"): + config.router_bias = False + if not hasattr(config, "router_dtype"): + config.router_dtype = "float32" + if not hasattr(config, "routed_scaling_factor"): + config.routed_scaling_factor = 1.0 + + return config + + +def _get_longcat_moe_quant_config( + config: _PretrainedConfig, + quant_config: _QuantizationConfig | None, + prefix: str, +): + if quant_config is None: + return None + + ignored_layers = quant_config.ignored_layers + if not ignored_layers: + return quant_config + + expert_proj_names = ("gate_proj", "up_proj", "down_proj") + num_expected = config.n_routed_experts * len(expert_proj_names) + num_ignored = 0 + for expert_id in range(config.n_routed_experts): + expert_prefix = add_prefix(f"experts.{expert_id}", prefix) + for proj_name in expert_proj_names: + if _should_ignore_quant_layer( + prefix=add_prefix(proj_name, expert_prefix), + ignored_layers=ignored_layers, + ): + num_ignored += 1 + + if num_ignored == 0: + return quant_config + if num_ignored == num_expected: + return None + + raise ValueError( + f"LongCat MoE layer {prefix} has partially ignored expert quantization " + f"({num_ignored}/{num_expected} expert projections). TokenSpeed requires " + "all experts in one fused MoE layer to use the same weight format." + ) + + +class _RuntimeLongcatRouter(nn.Module): + def __init__(self, config: _PretrainedConfig, prefix: str = ""): + super().__init__() + if getattr(config, "router_bias", False): + raise ValueError("LongCat router bias is not supported.") + + num_logits = config.n_routed_experts + config.zero_expert_num + params_dtype = ( + torch.bfloat16 if config.router_dtype == "bfloat16" else torch.float32 + ) + self.classifier = ReplicatedLinear( + config.hidden_size, + num_logits, + bias=False, + params_dtype=params_dtype, + quant_config=None, + prefix=add_prefix("classifier", prefix), + ) + self.e_score_correction_bias = nn.Parameter( + torch.zeros(num_logits, dtype=torch.float32) + ) + + def forward(self, hidden_states: torch.Tensor): + if _longcat_is_hopper_plus and hidden_states.shape[0] > 0: + return _dsv3_router_gemm( + hidden_states, + self.classifier.weight, + out_dtype=torch.float32, + enable_pdl=_pdl_enabled(), + ) + return _F.linear(hidden_states.float(), self.classifier.weight.float(), None) + + +class _RuntimeLongcatMoE(nn.Module): + def __init__( + self, + config: _PretrainedConfig, + mapping: _Mapping, + quant_config: _QuantizationConfig | None = None, + layer_index: int = -1, + prefix: str = "", + alt_stream: torch.cuda.Stream | None = None, + ): + super().__init__() + self.mapping = mapping + self.layer_index = layer_index + self.n_routed_experts = config.n_routed_experts + self.zero_expert_num = config.zero_expert_num + self.zero_expert_type = config.zero_expert_type + self.routed_scaling_factor = config.routed_scaling_factor + self.stream_fork = _StreamFork(alt_stream) + + if self.mapping.moe.ep_size > config.n_routed_experts: + raise ValueError( + f"EP size {self.mapping.moe.ep_size} is greater than the number " + f"of LongCat routed experts {config.n_routed_experts}." + ) + if config.hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {config.hidden_act}. " + "Only silu is supported for LongCat." + ) + + self.router = _RuntimeLongcatRouter( + config=config, + prefix=add_prefix("router", prefix), + ) + self.experts = _MoELayer( + top_k=config.moe_topk, + num_experts=( + config.n_routed_experts + + global_server_args_dict["ep_num_redundant_experts"] + ), + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + quant_config=quant_config, + layer_index=layer_index, + prefix=prefix, + tp_rank=self.mapping.moe.tp_rank, + tp_size=self.mapping.moe.tp_size, + ep_rank=self.mapping.moe.ep_rank, + ep_size=self.mapping.moe.ep_size, + zero_expert_type=config.zero_expert_type, + routing_config={ + "routed_scaling_factor": self.routed_scaling_factor, + "normalize_topk_weights": config.norm_topk_prob, + "correction_bias": self.router.e_score_correction_bias[ + : config.n_routed_experts + ], + "routing_method_type": _RoutingMethodType.DeepSeekV3, + }, + ) + if config.zero_expert_num > 0 and self.experts.topk_output_format.is_bypassed(): + raise ValueError( + "LongCat zero experts require a MoE backend that accepts " + "precomputed top-k ids. Launch with --moe-runner-backend triton." + ) + self.topk = _TopK( + top_k=config.moe_topk, + renormalize=config.norm_topk_prob, + correction_bias=self.router.e_score_correction_bias, + routed_scaling_factor=self.routed_scaling_factor, + output_format=_TopKOutputFormat.STANDARD, + zero_expert_num=config.zero_expert_num, + topk_indices_dtype=( + torch.int64 + if global_server_args_dict.get("enable_deep_ep", False) + else torch.int32 + ), + ) + + def get_moe_routed_weights(self): + return [ + param.data + for name, param in self.experts.named_parameters() + if name not in ["correction_bias"] and "shared_experts" not in name + ] + + def _apply_zero_experts(self, hidden_states: torch.Tensor, topk_output): + if self.zero_expert_num <= 0: + return None + + zero_expert_mask = (topk_output.topk_ids < 0) | ( + topk_output.topk_ids >= self.n_routed_experts + ) + zero_expert_weights = torch.where( + zero_expert_mask, + topk_output.topk_weights, + torch.zeros_like(topk_output.topk_weights), + ) + # Fused MoE kernels still read every selected expert id while building + # the dispatch plan, so zero-expert slots must keep a valid id. + topk_output.topk_ids[zero_expert_mask] = 0 + topk_output.topk_weights[zero_expert_mask] = 0.0 + + if self.zero_expert_type in ("identity", "copy"): + zero_weight = zero_expert_weights.sum(dim=-1, keepdim=True).to( + hidden_states.dtype + ) + return hidden_states * zero_weight + if self.zero_expert_type in ("", "drop"): + return None + raise ValueError( + f"Unsupported LongCat zero expert type: {self.zero_expert_type}" + ) + + def forward( + self, + hidden_states: torch.Tensor, + num_global_tokens: int, + max_num_tokens_per_gpu: int, + ) -> torch.Tensor: + with self.stream_fork.scope(enable=_get_is_capture_mode()): + router_logits = self.router(hidden_states) + if hidden_states.shape[0] > 0: + topk_output = self.topk(hidden_states, router_logits) + else: + topk_output = self.topk.empty_topk_output( + hidden_states.device, + hidden_states=hidden_states, + router_logits=router_logits, + ) + + zero_expert_output = self._apply_zero_experts(hidden_states, topk_output) + deferred_finalize = self.experts.supports_deferred_finalize + routed_expert_output = self.experts( + hidden_states=hidden_states, + topk_output=topk_output, + num_global_tokens=num_global_tokens, + max_num_tokens_per_gpu=max_num_tokens_per_gpu, + do_finalize=not deferred_finalize, + ) + + if deferred_finalize: + gemm2_out, expert_weights, expanded_idx = routed_expert_output + return _moe_finalize_fuse_shared( + gemm2_out, + expanded_idx, + expert_weights, + zero_expert_output, + top_k=self.topk.topk_config.top_k, + enable_pdl=_pdl_enabled(), + ) + + if zero_expert_output is not None: + routed_expert_output = routed_expert_output + zero_expert_output + return routed_expert_output + + +class _RuntimeLongcatDecoderLayer(nn.Module): + def __init__( + self, + config: _PretrainedConfig, + layer_id: int, + mapping: _Mapping, + quant_config: _QuantizationConfig | None = None, + prefix: str = "", + alt_stream: torch.cuda.Stream | None = None, + ) -> None: + super().__init__() + self.mapping = mapping + self.layer_id = layer_id + self.hidden_size = config.hidden_size + + rope_theta = _get_rope_theta(config) + rope_scaling = getattr(config, "rope_scaling", None) + if rope_scaling and "factor" not in rope_scaling: + rope_scaling = None + max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + + self.self_attn = nn.ModuleList( + [ + _DeepseekV3AttentionMLA( + config=config, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + qk_nope_head_dim=config.qk_nope_head_dim, + qk_rope_head_dim=config.qk_rope_head_dim, + v_head_dim=config.v_head_dim, + q_lora_rank=getattr(config, "q_lora_rank", None), + kv_lora_rank=config.kv_lora_rank, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + max_position_embeddings=max_position_embeddings, + quant_config=( + None + if "self_attn" in getattr(config, "disable_quant_module", []) + else quant_config + ), + layer_id=layer_id * 2 + branch_id, + prefix=add_prefix(f"self_attn.{branch_id}", prefix), + reduce_attn_results=False, + alt_stream=alt_stream, + mapping=self.mapping, + ) + for branch_id in range(2) + ] + ) + self.input_layernorm = nn.ModuleList( + [_RMSNorm(config.hidden_size, eps=config.rms_norm_eps) for _ in range(2)] + ) + self.post_attention_layernorm = nn.ModuleList( + [_RMSNorm(config.hidden_size, eps=config.rms_norm_eps) for _ in range(2)] + ) + dense_quant_config = ( + None + if "mlps" in getattr(config, "disable_quant_module", []) + else quant_config + ) + self.mlps = nn.ModuleList( + [ + _DeepseekV3MLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + mapping=self.mapping, + quant_config=dense_quant_config, + prefix=add_prefix(f"mlps.{branch_id}", prefix), + is_shared_expert=False, + ) + for branch_id in range(2) + ] + ) + self.mlp = _RuntimeLongcatMoE( + config=config, + mapping=self.mapping, + quant_config=_get_longcat_moe_quant_config( + config, + quant_config, + add_prefix("mlp", prefix), + ), + layer_index=layer_id, + prefix=add_prefix("mlp", prefix), + alt_stream=alt_stream, + ) + + self.moe_comm = _CommManager( + mapping=self.mapping, + layer_id=self.layer_id, + is_moe=True, + prev_is_moe=False, + input_layernorm=self.input_layernorm[0], + post_attn_layernorm=self.post_attention_layernorm[0], + ) + self.branch_comm = [ + _CommManager( + mapping=self.mapping, + layer_id=self.layer_id * 2 + branch_id, + is_moe=False, + prev_is_moe=False, + input_layernorm=self.input_layernorm[branch_id], + post_attn_layernorm=self.post_attention_layernorm[branch_id], + ) + for branch_id in range(2) + ] + self.final_norm_comm = self.branch_comm[1] + + def _forward_dense_mlp( + self, + branch_id: int, + hidden_states: torch.Tensor, + residual: torch.Tensor, + ctx: _ForwardContext, + ): + comm = self.branch_comm[branch_id] + hidden_states = comm.pre_mlp_comm(hidden_states, ctx) + hidden_states = self.mlps[branch_id](hidden_states) + hidden_states, residual = comm.post_mlp_fused(hidden_states, residual, ctx) + return hidden_states, residual + + def _forward_moe( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor, + ctx: _ForwardContext, + num_global_tokens: int, + max_num_tokens_per_gpu: int, + ): + hidden_states = self.moe_comm.pre_mlp_comm(hidden_states, ctx) + hidden_states = self.mlp( + hidden_states, + num_global_tokens, + max_num_tokens_per_gpu, + ) + hidden_states, residual = self.moe_comm.post_mlp_fused( + hidden_states, + residual, + ctx, + ) + return hidden_states, residual + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: _ForwardContext, + out_cache_loc: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + num_global_tokens, max_num_tokens_per_gpu = self.moe_comm.get_num_tokens(ctx) + + if ctx.forward_mode.is_idle(): + hidden_states, residual = self._forward_moe( + hidden_states, + residual, + ctx, + num_global_tokens, + max_num_tokens_per_gpu, + ) + return hidden_states, residual + + hidden_states, residual = self.moe_comm.input_reduce_norm( + hidden_states, + residual, + ) + hidden_states = self.self_attn[0]( + positions=positions, + hidden_states=hidden_states, + ctx=ctx, + out_cache_loc=out_cache_loc, + comm_manager=self.moe_comm, + ) + hidden_states, residual = self.moe_comm.post_attn_reduce_norm( + hidden_states, + residual, + ctx, + ) + + branch_input = hidden_states + branch_residual = residual + moe_hidden_states, _ = self._forward_moe( + branch_input, + branch_residual, + ctx, + num_global_tokens, + max_num_tokens_per_gpu, + ) + + hidden_states, residual = self._forward_dense_mlp( + 0, + branch_input, + branch_residual, + ctx, + ) + hidden_states, residual = self.branch_comm[1].input_reduce_norm( + hidden_states, + residual, + ) + hidden_states = self.self_attn[1]( + positions=positions, + hidden_states=hidden_states, + ctx=ctx, + out_cache_loc=out_cache_loc, + comm_manager=self.branch_comm[1], + ) + hidden_states, residual = self.branch_comm[1].post_attn_reduce_norm( + hidden_states, + residual, + ctx, + ) + hidden_states, residual = self._forward_dense_mlp( + 1, + hidden_states, + residual, + ctx, + ) + + hidden_states = hidden_states + moe_hidden_states + return hidden_states, residual + + +class _RuntimeLongcatModel(nn.Module): + fall_back_to_pt_during_load = False + + def __init__( + self, + config: _PretrainedConfig, + mapping: _Mapping, + quant_config: _QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + _ensure_longcat_config(config) + self.mapping = mapping + self.padding_id = getattr(config, "pad_token_id", None) + self.vocab_size = config.vocab_size + + self.embed_tokens = _VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + self.alt_stream = torch.cuda.Stream() if torch.cuda.is_available() else None + self.layers = nn.ModuleList( + [ + _RuntimeLongcatDecoderLayer( + config, + layer_id, + mapping=self.mapping, + quant_config=quant_config, + prefix=add_prefix(f"layers.{layer_id}", prefix), + alt_stream=self.alt_stream, + ) + for layer_id in range(config.num_hidden_layers) + ] + ) + self.norm = _RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.layers_to_capture: set[int] = set() + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + ctx: _ForwardContext, + out_cache_loc: torch.Tensor, + input_embeds: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, list[torch.Tensor] | None]: + if input_embeds is not None: + hidden_states = input_embeds + else: + hidden_states = self.embed_tokens(input_ids) + + residual = None + aux_hidden_states = [] if self.layers_to_capture else None + layer = None + for layer_id, layer in enumerate(self.layers): + if aux_hidden_states is not None and layer_id in self.layers_to_capture: + aux_hidden_states.append( + hidden_states + residual if residual is not None else hidden_states + ) + with _get_global_expert_distribution_recorder().with_current_layer( + layer_id + ): + hidden_states, residual = layer( + positions, + hidden_states, + ctx, + out_cache_loc, + residual, + ) + + if not ctx.forward_mode.is_idle() and layer is not None: + hidden_states, _ = layer.final_norm_comm.final_norm( + hidden_states, + residual, + ctx, + self.norm, + ) + return hidden_states, aux_hidden_states + + +class LongcatFlashForCausalLM(_BaseCausalLM): + model_cls = _RuntimeLongcatModel + + def __init__( + self, + config: _PretrainedConfig, + mapping: _Mapping, + model: _RuntimeLongcatModel | None = None, + quant_config: _QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + _ensure_longcat_config(config) + self._model_override = model + super().__init__( + config=config, + mapping=mapping, + quant_config=quant_config, + prefix=prefix, + ) + + def resolve_model( + self, + config: _PretrainedConfig, + mapping: _Mapping, + quant_config: _QuantizationConfig | None, + prefix: str, + ) -> _RuntimeLongcatModel: + if self._model_override is not None: + return self._model_override + return self.model_cls( + config, + mapping=mapping, + quant_config=quant_config, + prefix=add_prefix("model", prefix), + ) + + def post_init(self) -> None: + self._routed_experts_weights_of_layer = LazyValue( + lambda: { + layer_id: layer.mlp.get_moe_routed_weights() + for layer_id, layer in enumerate(self.model.layers) + if isinstance(layer.mlp, _RuntimeLongcatMoE) + } + ) + + @property + def routed_experts_weights_of_layer(self): + return self._routed_experts_weights_of_layer.value + + def set_eagle3_layers_to_capture(self, layer_ids: list[int] | None = None): + self.capture_aux_hidden_states = True + if layer_ids is None: + num_layers = self.config.num_hidden_layers + self.model.layers_to_capture = {2, num_layers // 2, num_layers - 3} + else: + self.model.layers_to_capture = {val + 1 for val in layer_ids} + + def get_param(self, params_dict, name): + if name in params_dict: + return params_dict[name] + if "language_model." in name: + name = name.replace("language_model.", "") + if name in params_dict: + return params_dict[name] + if ".mtp." in name or name.startswith("model.mtp."): + return None + if name.endswith(_LONGCAT_OPTIONAL_MISSING_WEIGHT_SUFFIXES): + return None + _longcat_logger.warning("The %s is not in the model.", name) + return None + + def load_weights(self, weights: _Iterable[tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + fuse_qkv_a_proj = getattr(self.config, "q_lora_rank", None) is not None + params_dict = dict(self.named_parameters()) + moe_loader = _build_moe_checkpoint_loader( + params_dict=params_dict, + expert_schema=_ExpertCheckpointSchema( + gate_proj_name="gate_proj", + down_proj_name="down_proj", + up_proj_name="up_proj", + ), + num_experts=self.config.n_routed_experts, + ep_rank=self.mapping.moe.ep_rank, + ep_size=self.mapping.moe.ep_size, + ) + + for name, loaded_weight in weights: + layer_id = _get_layer_id(name) + if ( + layer_id is not None + and hasattr(self.model, "start_layer") + and ( + layer_id < self.model.start_layer + or layer_id >= self.model.end_layer + ) + ): + continue + if "rotary_emb.inv_freq" in name: + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if "mlp.experts." in name and name not in params_dict: + continue + mapped_name = name.replace(weight_name, param_name) + if mapped_name.endswith(".bias") and mapped_name not in params_dict: + continue + param = self.get_param(params_dict, mapped_name) + if param is None: + break + param.weight_loader(param, loaded_weight, shard_id) + break + else: + if name.endswith(".bias") and name not in params_dict: + continue + if moe_loader.matches(name): + moe_loader.load(name, loaded_weight) + continue + + if fuse_qkv_a_proj and ( + "q_a_proj" in name or "kv_a_proj_with_mqa" in name + ): + quant_block_size = 1 + if ( + self.quant_config is not None + and self.quant_config.weight_block_size is not None + ): + quant_block_size = self.quant_config.weight_block_size[0] + begin_size_by_name = { + "q_a_proj": 0, + "kv_a_proj_with_mqa": self.config.q_lora_rank, + } + if "q_a_proj" in name: + param = self.get_param( + params_dict, + name.replace("q_a_proj", "fused_qkv_a_proj_with_mqa"), + ) + begin_size = begin_size_by_name["q_a_proj"] + else: + param = self.get_param( + params_dict, + name.replace( + "kv_a_proj_with_mqa", + "fused_qkv_a_proj_with_mqa", + ), + ) + begin_size = begin_size_by_name["kv_a_proj_with_mqa"] + if param is None: + continue + if "scale_inv" in name: + begin_size //= quant_block_size + param.weight_loader(param, loaded_weight, begin_size=begin_size) + continue + + if "q_a_proj" in name and name not in params_dict: + name = name.replace("q_a_proj", "q_proj") + param = self.get_param(params_dict, name) + if param is None: + continue + weight_loader = getattr(param, "weight_loader", _default_weight_loader) + weight_loader(param, loaded_weight) + + self.post_load_weights() + + def post_load_weights(self): + for layer in self.model.layers: + for self_attn in layer.self_attn: + if hasattr( + self.quant_config, "weight_block_size" + ) and self_attn.kv_b_proj.weight.dtype in ( + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + ): + weight_block_size = self.quant_config.weight_block_size + if weight_block_size is not None: + if not hasattr(self_attn.kv_b_proj, "weight_scale_inv"): + raise RuntimeError( + "kv_b_proj.weight_scale_inv is required for block FP8 dequant." + ) + dtype = torch.get_default_dtype() + w = _block_dequant( + self_attn.kv_b_proj.weight, + self_attn.kv_b_proj.weight_scale_inv, + weight_block_size, + ).to(dtype) + else: + w = self_attn.kv_b_proj.weight + else: + w = self_attn.kv_b_proj.weight + + 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) + self_attn.w_kc = w_kc.transpose(1, 2).contiguous().transpose(1, 2) + self_attn.w_vc = w_vc.contiguous().transpose(1, 2) + if getattr(self.config, "mla_scale_q_lora", False) and hasattr( + self_attn, + "q_a_layernorm", + ): + self_attn.q_a_layernorm.weight.data *= ( + self.config.hidden_size / self.config.q_lora_rank + ) ** 0.5 + if getattr(self.config, "mla_scale_kv_lora", False): + self_attn.kv_a_layernorm.weight.data *= ( + self.config.hidden_size / self.config.kv_lora_rank + ) ** 0.5 + + def load_kv_cache_scales(self, quantization_param_path: str) -> None: + tp_size = self.mapping.attn.tp_size + tp_rank = self.mapping.attn.tp_rank + for attn_idx, scaling_factor in _kv_cache_scales_loader( + quantization_param_path, + tp_rank, + tp_size, + self.config.num_hidden_layers * 2, + self.config.__class__.model_type, + ): + layer_idx, branch_idx = divmod(attn_idx, 2) + if not isinstance(self.model.layers[layer_idx], nn.Identity): + self_attn = self.model.layers[layer_idx].self_attn[branch_idx] + for attn in (self_attn.attn_mha, self_attn.attn_mqa): + if attn is not None and hasattr(attn, "k_scale"): + attn.k_scale = scaling_factor + attn.k_scale_float = scaling_factor + + def get_embed_and_head(self): + return self.model.embed_tokens.weight, self.lm_head.weight + + def set_embed_and_head(self, embed, head): + del self.model.embed_tokens.weight + del self.lm_head.weight + self.model.embed_tokens.weight = embed + self.lm_head.weight = head + torch.cuda.empty_cache() + torch.cuda.synchronize() + + @classmethod + def get_model_config_for_expert_location(cls, config): + _ensure_longcat_config(config) + return _ModelConfigForExpertLocation( + num_layers=config.num_hidden_layers, + num_logical_experts=config.n_routed_experts, + num_groups=None, + ) + + +FLASHForCausalLM = LongcatFlashForCausalLM +EntryClass = LongcatFlashForCausalLM diff --git a/python/tokenspeed/runtime/models/minimax_m2.py b/python/tokenspeed/runtime/models/minimax_m2.py new file mode 100644 index 0000000..e628901 --- /dev/null +++ b/python/tokenspeed/runtime/models/minimax_m2.py @@ -0,0 +1,933 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Inference-only MiniMax-M2 family model compatible with HuggingFace weights.""" + +# ruff: noqa: E402 + +import logging +from collections.abc import Iterable +from typing import Any, cast + +import torch +import triton +import triton.language as tl +from tokenspeed_kernel.ops.communication.trtllm import ( + minimax_allreduce_rms_qk, + trtllm_create_ipc_workspace_for_minimax, +) +from tokenspeed_kernel.platform import current_platform +from tokenspeed_kernel.torch_compile import get_compiler_backend +from torch import nn + +from tokenspeed.runtime.configs.minimax_m2_config import MiniMaxM2Config +from tokenspeed.runtime.distributed.comm_ops import all_reduce +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.layers.linear import ( + QKVParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from tokenspeed.runtime.layers.logits_processor import LogitsProcessor +from tokenspeed.runtime.layers.moe import ( + ExpertCheckpointSchema, + build_moe_checkpoint_loader, +) +from tokenspeed.runtime.layers.moe.topk import TopK +from tokenspeed.runtime.layers.moe.utils import RoutingMethodType +from tokenspeed.runtime.layers.paged_attention import PagedAttention +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.layers.rotary_embedding import get_rope +from tokenspeed.runtime.layers.vocab_parallel_embedding import ParallelLMHead +from tokenspeed.runtime.model_loader.weight_utils import ( + default_weight_loader, + sharded_weight_loader, +) +from tokenspeed.runtime.models.base import ( + BaseCausalLM, + BaseMoEDecoderLayer, + BaseTransformerModel, +) +from tokenspeed.runtime.models.utils import ( + create_fused_set_kv_buffer_arg, + validate_attention_partition, +) +from tokenspeed.runtime.moe.expert_location import ModelConfigForExpertLocation +from tokenspeed.runtime.utils import ( + LazyValue, + add_prefix, + set_weight_attrs, +) +from tokenspeed.runtime.utils.env import envs, global_server_args_dict +from tokenspeed.runtime.utils.pdl import pdl_enabled + +logger = logging.getLogger(__name__) + +_is_nvidia = current_platform().is_nvidia + +if _is_nvidia: + from tokenspeed_kernel.thirdparty.cuda import fp32_router_gemm + +from tokenspeed.runtime.layers.moe.expert import MoELayer as _MoELayer + +MoELayer = _MoELayer + + +class MiniMaxM2SparseMoeBlock(nn.Module): + def __init__( + self, + config: MiniMaxM2Config, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + layer_index: int = -1, + prefix: str = "", + ): + super().__init__() + + self.mapping = mapping + self.layer_index = layer_index + self.tp_size = mapping.world_size + + if self.tp_size > config.num_local_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than " + f"the number of experts {config.num_local_experts}." + ) + + self.gate = ReplicatedLinear( + config.hidden_size, + config.num_local_experts, + bias=False, + quant_config=None, + params_dtype=torch.float32, + prefix=add_prefix("gate", prefix), + ) + + if config.use_routing_bias: + self.routing_bias = nn.Parameter( + torch.zeros(config.num_local_experts, dtype=torch.float32) + ) + else: + self.routing_bias = None + + self.use_fp32_router_gemm = ( + current_platform().is_hopper_plus + and config.hidden_size == 3072 + and config.num_local_experts == 256 + ) + + routing_config = { + "n_group": 1, + "topk_group": 1, + "routed_scaling_factor": 1.0, + "normalize_topk_weights": True, + "correction_bias": self.routing_bias, + "routing_method_type": RoutingMethodType.MiniMax2, + } + + self.experts = MoELayer( + top_k=config.num_experts_per_tok, + num_experts=config.num_local_experts + + global_server_args_dict["ep_num_redundant_experts"], + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + quant_config=quant_config, + layer_index=layer_index, + prefix=prefix, + tp_rank=self.mapping.moe.tp_rank, + tp_size=self.mapping.moe.tp_size, + ep_rank=self.mapping.moe.ep_rank, + ep_size=self.mapping.moe.ep_size, + routing_config=routing_config, + ) + self.topk = TopK( + top_k=config.num_experts_per_tok, + renormalize=True, + use_grouped_topk=True, + num_expert_group=1, + topk_group=1, + correction_bias=self.routing_bias, + routed_scaling_factor=1.0, + output_format=self.experts.topk_output_format, + ) + + def get_moe_routed_weights(self): + + return [ + x.data + for name, x in self.experts.named_parameters() + if name not in ["correction_bias"] + ] + + def forward( + self, + hidden_states: torch.Tensor, + num_global_tokens: int, + max_num_tokens_per_gpu: int, + ) -> torch.Tensor: + + num_tokens, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + + # FP32 Router GEMM. + if self.use_fp32_router_gemm and hidden_states.shape[0] > 0: + router_logits = fp32_router_gemm(hidden_states, self.gate.weight) + else: + router_logits, _ = self.gate(hidden_states.to(torch.float32)) + + if hidden_states.shape[0] > 0: + topk_output = self.topk(hidden_states, router_logits) + else: + topk_output = self.topk.empty_topk_output( + hidden_states.device, + hidden_states=hidden_states, + router_logits=router_logits, + ) + + # Experts. + final_hidden_states = self.experts( + hidden_states=hidden_states, + topk_output=topk_output, + num_global_tokens=num_global_tokens, + max_num_tokens_per_gpu=max_num_tokens_per_gpu, + ) + + return final_hidden_states.view(num_tokens, hidden_dim) + + +@triton.jit +def _rmsnorm_sumsq_kernel( + x1_ptr, + x2_ptr, + stride_x1, + stride_x2, + sum_sq_ptr, + B, + D1, + D2, + BLOCK_SIZE1: tl.constexpr, + BLOCK_SIZE2: tl.constexpr, +): + row_id = tl.program_id(0) + x1_row = x1_ptr + row_id * stride_x1 + x2_row = x2_ptr + row_id * stride_x2 + + offsets1 = tl.arange(0, BLOCK_SIZE1) + offsets2 = tl.arange(0, BLOCK_SIZE2) + + x1 = tl.load(x1_row + offsets1, mask=offsets1 < D1, other=0.0).to(tl.float32) + x2 = tl.load(x2_row + offsets2, mask=offsets2 < D2, other=0.0).to(tl.float32) + + tl.store(sum_sq_ptr + row_id, tl.sum(x1 * x1, axis=0)) + tl.store(sum_sq_ptr + row_id + B, tl.sum(x2 * x2, axis=0)) + + +@triton.jit +def _rmsnorm_apply_kernel( + x1_ptr, + x2_ptr, + w1_ptr, + w2_ptr, + sum_sq_ptr, + out1_ptr, + out2_ptr, + B, + D1, + D2, + stride_x1, + stride_x2, + tp_world, + eps, + BLOCK_SIZE1: tl.constexpr, + BLOCK_SIZE2: tl.constexpr, +): + row_id = tl.program_id(0) + x1_row = x1_ptr + row_id * stride_x1 + x2_row = x2_ptr + row_id * stride_x2 + out1_row = out1_ptr + row_id * stride_x1 + out2_row = out2_ptr + row_id * stride_x2 + + inv_rms1 = tl.rsqrt(tl.load(sum_sq_ptr + row_id) / D1 / tp_world + eps) + inv_rms2 = tl.rsqrt(tl.load(sum_sq_ptr + row_id + B) / D2 / tp_world + eps) + + offsets1 = tl.arange(0, BLOCK_SIZE1) + offsets2 = tl.arange(0, BLOCK_SIZE2) + mask1 = offsets1 < D1 + mask2 = offsets2 < D2 + + x1 = tl.load(x1_row + offsets1, mask=mask1, other=0.0) + w1 = tl.load(w1_ptr + offsets1, mask=mask1, other=1.0) + x2 = tl.load(x2_row + offsets2, mask=mask2, other=0.0) + w2 = tl.load(w2_ptr + offsets2, mask=mask2, other=1.0) + + tl.store( + out1_row + offsets1, + (x1.to(tl.float32) * inv_rms1 * w1.to(tl.float32)).to(x1.dtype), + mask=mask1, + ) + tl.store( + out2_row + offsets2, + (x2.to(tl.float32) * inv_rms2 * w2.to(tl.float32)).to(x2.dtype), + mask=mask2, + ) + + +@torch.compile(dynamic=True, backend=get_compiler_backend()) +def fused_qk_rmsnorm_triton( + q: torch.Tensor, + k: torch.Tensor, + q_weight: torch.Tensor, + k_weight: torch.Tensor, + tp_size: int, + tp_rank: int, + tp_group: tuple[int, ...], + eps: float = 1e-6, +) -> tuple[torch.Tensor, torch.Tensor]: + """Fused QK RMSNorm: sumsq → allreduce → apply, using 2 Triton kernels.""" + q = q.contiguous() + k = k.contiguous() + + B, D1 = q.shape + _, D2 = k.shape + BLOCK_SIZE1 = triton.next_power_of_2(D1) + BLOCK_SIZE2 = triton.next_power_of_2(D2) + + # Pad for allreduce alignment (16-byte = 4 floats) + B_padded = (B + B + 3) // 4 * 4 + sum_sq = torch.empty(B_padded, device=q.device, dtype=torch.float32) + + _rmsnorm_sumsq_kernel[(B,)]( + q, + k, + q.stride(0), + k.stride(0), + sum_sq, + B, + D1, + D2, + BLOCK_SIZE1, + BLOCK_SIZE2, + ) + + if tp_size > 1: + sum_sq = all_reduce(sum_sq, tp_group) + + out1 = torch.empty_like(q) + out2 = torch.empty_like(k) + + _rmsnorm_apply_kernel[(B,)]( + q, + k, + q_weight, + k_weight, + sum_sq, + out1, + out2, + B, + D1, + D2, + q.stride(0), + k.stride(0), + tp_size, + eps, + BLOCK_SIZE1, + BLOCK_SIZE2, + ) + + return out1, out2 + + +def _minimax_fast_path_available( + q: torch.Tensor, + k: torch.Tensor, + tp_size: int, +) -> bool: + """Fast-path CUDA kernel (Lamport AR fused with RMSNorm) is usable only for + TP in {2,4,8,16} and global head dims (Q, K) == (6144, 1024).""" + if tp_size not in (2, 4, 8, 16): + return False + if q.dim() != 2 or k.dim() != 2: + return False + if q.shape[-1] * tp_size != 6144 or k.shape[-1] * tp_size != 1024: + return False + if q.dtype not in (torch.float16, torch.bfloat16): + return False + return True + + +class _MinimaxARWorkspace: + """Singleton holder for the dedicated MiniMax AR+RMSNorm IPC workspace. + + One workspace per (tp_group, dtype_elem_size, max_token_num). Lifetime is + tied to the process; it lives as long as the model. + """ + + def __init__(self) -> None: + self._entries: dict[tuple[tuple[int, ...], int, int], dict[str, Any]] = {} + + def get_or_create( + self, + tp_rank: int, + tp_group: tuple[int, ...], + max_token_num: int, + dtype_elem_size: int, + ) -> torch.Tensor | None: + key = (tp_group, dtype_elem_size, max_token_num) + # Grow max_token_num if needed: find any existing entry for the same + # (group, dtype) and check whether we can reuse it. + for (g, sz, cap), entry in self._entries.items(): + if g == tp_group and sz == dtype_elem_size and cap >= max_token_num: + return entry["workspace"] + + from tokenspeed.runtime.distributed.process_group_manager import ( + process_group_manager as pg_manager, + ) + + device_group = pg_manager.get_process_group("nccl", tp_group) + try: + ipc_handles, workspace = trtllm_create_ipc_workspace_for_minimax( + tp_rank=tp_rank, + tp_size=len(tp_group), + max_token_num=max_token_num, + group=device_group, + dtype_elem_size=dtype_elem_size, + ) + except Exception: + logger.exception("Failed to create MiniMax AR+RMSNorm IPC workspace") + return None + + self._entries[key] = { + "ipc_handles": ipc_handles, + "workspace": workspace, + "device_group": device_group, + } + return workspace + + +_minimax_ar_workspace = _MinimaxARWorkspace() + + +_FORCE_TRITON_AR_RMSNORM = envs.TOKENSPEED_MINIMAX_AR_USE_TRITON.get() + + +def fused_qk_rmsnorm( + q: torch.Tensor, + k: torch.Tensor, + q_weight_fp32: torch.Tensor, + k_weight_fp32: torch.Tensor, + q_weight_bf16: torch.Tensor | None, + k_weight_bf16: torch.Tensor | None, + tp_size: int, + tp_rank: int, + tp_group: tuple[int, ...], + eps: float = 1e-6, +) -> tuple[torch.Tensor, torch.Tensor]: + """Route to the Lamport fused-AR QK RMSNorm kernel when its shape + constraints hold, else fall back to the Triton sumsq/apply path. + Setting TOKENSPEED_MINIMAX_AR_USE_TRITON=1 forces the Triton path (A/B debug).""" + if ( + not _FORCE_TRITON_AR_RMSNORM + and q_weight_bf16 is not None + and k_weight_bf16 is not None + and _minimax_fast_path_available(q, k, tp_size) + ): + num_tokens = q.shape[0] + # Allocate once with a generous ceiling so batch-size changes never + # force reallocation. 16384 tokens × TP=16 fits in ~6MB of lamport buffer. + _MINIMAX_WORKSPACE_CAP = 16384 + workspace = _minimax_ar_workspace.get_or_create( + tp_rank=tp_rank, + tp_group=tp_group, + max_token_num=max(num_tokens, _MINIMAX_WORKSPACE_CAP), + dtype_elem_size=q.element_size(), + ) + if workspace is not None: + # Kernel reads q/k at their row stride (q_row_stride_f4), so a + # non-contiguous slice from a fused-QKV split is fine. + return minimax_allreduce_rms_qk( + q=q, + k=k, + norm_weight_q=q_weight_bf16, + norm_weight_k=k_weight_bf16, + workspace_ptrs=workspace, + rank=tp_rank, + nranks=tp_size, + eps=eps, + trigger_completion_at_end=True, + launch_with_pdl=pdl_enabled(), + ) + + return fused_qk_rmsnorm_triton( + q, k, q_weight_fp32, k_weight_fp32, tp_size, tp_rank, tp_group, eps + ) + + +class MiniMaxM2RMSNormTP(nn.Module): + """Tensor-parallel RMSNorm for MiniMax Q/K normalization.""" + + def __init__( + self, + global_hidden_size: int, + tp_rank: int, + tp_size: int, + tp_group: tuple[int, ...], + eps: float = 1e-6, + ) -> None: + super().__init__() + if global_hidden_size % tp_size != 0: + raise ValueError( + f"global_hidden_size={global_hidden_size} must be divisible by tp_size={tp_size}." + ) + self.local_hidden_size = global_hidden_size // tp_size + self.tp_rank = tp_rank + self.tp_size = tp_size + self.tp_group = tp_group + self.variance_epsilon = eps + + self.weight = nn.Parameter(torch.ones(self.local_hidden_size)) + self._weight_bf16: torch.Tensor | None = None + self._weight_bf16_src_ptr: int = 0 + + set_weight_attrs( + self.weight, {"weight_loader": sharded_weight_loader(0, self.tp_rank)} + ) + + def bf16_weight(self) -> torch.Tensor: + # The Lamport-fused AR kernel requires bf16 gamma. Cache a bf16 copy + # of the fp32 Parameter; refresh if the backing storage ever changes + # (e.g. weights reloaded). + src_ptr = self.weight.data_ptr() + if self._weight_bf16 is None or self._weight_bf16_src_ptr != src_ptr: + self._weight_bf16 = self.weight.detach().to(torch.bfloat16).contiguous() + self._weight_bf16_src_ptr = src_ptr + return self._weight_bf16 + + +def remap_minimax_weight_name(name: str) -> str: + """Map HF checkpoint-only MiniMax names to local parameter names.""" + if "e_score_correction_bias" in name: + name = name.replace("e_score_correction_bias", "routing_bias") + if "block_sparse_moe" in name: + name = name.replace("block_sparse_moe", "mlp") + return name + + +def get_spec_layer_idx_from_weight_name( + config: MiniMaxM2Config, weight_name: str +) -> int | None: + """Return the extra speculative layer index encoded after main layers. + + Public MiniMax-M2 configs can carry speculative-decoding metadata even when + the released checkpoints do not include those extra layer weights. The + serving model instantiated here is main-model only, so extra layers beyond + ``num_hidden_layers`` should be ignored if a checkpoint ever includes them. + """ + num_spec_modules = int(getattr(config, "num_mtp_modules", 0) or 0) + layers_per_spec_module = int(getattr(config, "mtp_transformer_layers", 1) or 1) + num_spec_layers = num_spec_modules * layers_per_spec_module + start_layer = int(config.num_hidden_layers) + for i in range(num_spec_layers): + layer_idx = start_layer + i + if weight_name.startswith(f"model.layers.{layer_idx}."): + return layer_idx + return None + + +class MiniMaxM2Attention(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + mapping: Mapping, + layer_id: int = 0, + rope_theta: float = 10000, + rope_scaling: dict[str, Any] | None = None, + max_position_embeddings: int = 8192, + head_dim: int | None = None, + rotary_dim: int | None = None, + rms_norm_eps: float = 1e-06, + attention_bias: bool = False, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.layer_id = layer_id + self.mapping = mapping + self.hidden_size = hidden_size + self.attn_tp_size = mapping.attn.tp_size + self.attn_tp_rank = mapping.attn.tp_rank + self.attn_tp_group = mapping.attn.tp_group + self.total_num_heads = num_heads + self.total_num_kv_heads = num_kv_heads + validate_attention_partition( + self.total_num_heads, + self.total_num_kv_heads, + self.attn_tp_size, + ) + self.num_heads = self.total_num_heads // self.attn_tp_size + self.num_kv_heads = max(1, self.total_num_kv_heads // self.attn_tp_size) + self.head_dim = head_dim or hidden_size // self.total_num_heads + self.rotary_dim = rotary_dim or self.head_dim + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=attention_bias, + quant_config=quant_config, + tp_rank=self.attn_tp_rank, + tp_size=self.attn_tp_size, + tp_group=self.attn_tp_group, + prefix=add_prefix("qkv_proj", prefix), + ) + + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=attention_bias, + quant_config=quant_config, + reduce_results=False, + tp_rank=self.attn_tp_rank, + tp_size=self.attn_tp_size, + tp_group=self.attn_tp_group, + prefix=add_prefix("o_proj", prefix), + ) + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.rotary_dim, + max_position=max_position_embeddings, + base=int(rope_theta), + rope_scaling=rope_scaling, + ) + + self.attn = PagedAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + ) + + self.q_norm = MiniMaxM2RMSNormTP( + self.total_num_heads * self.head_dim, + tp_rank=self.attn_tp_rank, + tp_size=self.attn_tp_size, + tp_group=self.attn_tp_group, + eps=rms_norm_eps, + ) + + self.k_norm = MiniMaxM2RMSNormTP( + self.total_num_kv_heads * self.head_dim, + tp_rank=self.attn_tp_rank, + tp_size=self.attn_tp_size, + tp_group=self.attn_tp_group, + eps=rms_norm_eps, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + ) -> torch.Tensor: + + if hidden_states.shape[0] == 0: + return hidden_states + + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + + q, k = fused_qk_rmsnorm( + q, + k, + self.q_norm.weight, + self.k_norm.weight, + self.q_norm.bf16_weight(), + self.k_norm.bf16_weight(), + self.q_norm.tp_size, + self.q_norm.tp_rank, + self.q_norm.tp_group, + self.q_norm.variance_epsilon, + ) + + fused_kv_arg = None + if ctx.attn_backend.support_kv_cache_prewrite(): + n = q.shape[0] + v_3d = v.view(n, self.num_kv_heads, self.head_dim) + fused_kv_arg = create_fused_set_kv_buffer_arg( + value=v_3d, + layer=self.attn, + out_cache_loc=out_cache_loc, + token_to_kv_pool=ctx.token_to_kv_pool, + ) + + if fused_kv_arg is not None: + q_rope = torch.empty((n, self.q_size), dtype=q.dtype, device=q.device) + q, k = self.rotary_emb( + positions, + q, + k, + fused_set_kv_buffer_arg=fused_kv_arg, + output_q_rope=q_rope, + enable_pdl=pdl_enabled(), + ) + attn_output = self.attn( + q_rope, + None, + None, + save_kv_cache=False, + ctx=ctx, + out_cache_loc=out_cache_loc, + ) + + else: + + q, k = self.rotary_emb(positions, q, k) + q = q.view(-1, self.num_heads, self.head_dim) + k = k.view(-1, self.num_kv_heads, self.head_dim) + v = v.view(-1, self.num_kv_heads, self.head_dim) + attn_output = self.attn(q, k, v, ctx=ctx, out_cache_loc=out_cache_loc) + + output, _ = self.o_proj(attn_output) + + return output + + +class MiniMaxM2DecoderLayer(BaseMoEDecoderLayer): + + def __init__( + self, + config: MiniMaxM2Config, + layer_id: int, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + self._config = config + self._mapping = mapping + self._quant_config = quant_config + + super().__init__( + config=config, + layer_id=layer_id, + mapping=mapping, + quant_config=quant_config, + prefix=prefix, + ) + + def resolve_attn(self, prefix: str) -> nn.Module: + + config = self._config + + return MiniMaxM2Attention( + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + mapping=self._mapping, + layer_id=self.layer_id, + rope_theta=config.rope_theta, + rope_scaling=getattr(config, "rope_scaling", None), + max_position_embeddings=config.max_position_embeddings, + head_dim=config.head_dim, + rotary_dim=config.rotary_dim, + rms_norm_eps=config.rms_norm_eps, + attention_bias=config.attention_bias, + quant_config=self._quant_config, + prefix=add_prefix("self_attn", prefix), + ) + + def resolve_mlp(self, prefix: str) -> nn.Module: + + return MiniMaxM2SparseMoeBlock( + config=self._config, + mapping=self._mapping, + quant_config=self._quant_config, + layer_index=self.layer_id, + prefix=add_prefix("block_sparse_moe", prefix), + ) + + +class MiniMaxM2Model(BaseTransformerModel): + + layer_cls = MiniMaxM2DecoderLayer + + +class MiniMaxM2ForCausalLM(BaseCausalLM): + + model_cls = MiniMaxM2Model + fall_back_to_pt_during_load = False + + def __init__( + self, + config: MiniMaxM2Config, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + + super().__init__(config, mapping, quant_config, prefix) + + self._routed_experts_weights_of_layer = LazyValue( + lambda: { + layer_id: cast( + MiniMaxM2DecoderLayer, self.model.layers[layer_id] + ).mlp.get_moe_routed_weights() + for layer_id in range(len(self.model.layers)) + } + ) + + @property + def routed_experts_weights_of_layer(self): + return self._routed_experts_weights_of_layer.value + + def resolve_lm_head(self, config, quant_config, prefix): + + if self.mapping.attn.has_dp: + return ReplicatedLinear( + config.hidden_size, + config.vocab_size, + bias=False, + prefix=add_prefix("lm_head", prefix), + ) + + return ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=add_prefix("lm_head", prefix), + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + + def resolve_logits_processor(self, config): + + return LogitsProcessor( + config, + skip_all_gather=self.mapping.attn.has_dp, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]], **kwargs): + + stacked_params_mapping = [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ] + + # Skip loading extra parameters for GPTQ/nvfp4 models. + ignore_suffixes = ( + ".bias", + "_bias", + ".k_scale", + "_k_scale", + ".v_scale", + "_v_scale", + ".weight_scale", + "_weight_scale", + ".weight_scale_2", + "_weight_scale_2", + ".input_scale", + "_input_scale", + ) + + loaded_params: set[str] = set() + params_dict = dict(self.named_parameters(remove_duplicate=False)) + moe_loader = build_moe_checkpoint_loader( + params_dict=params_dict, + expert_schema=ExpertCheckpointSchema( + gate_proj_name="w1", + down_proj_name="w2", + up_proj_name="w3", + ), + num_experts=self.config.num_local_experts, + ep_rank=self.mapping.moe.ep_rank, + ep_size=self.mapping.moe.ep_size, + ) + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + if get_spec_layer_idx_from_weight_name(self.config, name) is not None: + continue + + name = remap_minimax_weight_name(name) + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if "mlp.experts." in name: + continue + + name = name.replace(weight_name, param_name) + if name.endswith(ignore_suffixes) and name not in params_dict: + continue + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # moe_loader.matches must be checked BEFORE the + # ignore_suffixes gate. Expert scale names end with + # `.weight_scale` / `.weight_scale_2` / `.input_scale` — those + # match ignore_suffixes, and the pre-remap checkpoint name + # (e.g. `experts.10.w1.weight_scale`) is not in params_dict, + # so the ignore gate would otherwise silently drop every FP4 + # expert scale and leave the layer with uninitialized scales. + if moe_loader.matches(name): + name = moe_loader.load(name, loaded_weight) + else: + if name.endswith(ignore_suffixes) and name not in params_dict: + continue + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params + + @classmethod + def get_model_config_for_expert_location(cls, config): + + return ModelConfigForExpertLocation( + num_layers=config.num_hidden_layers, + num_logical_experts=config.num_local_experts, + num_groups=None, + ) + + +EntryClass = MiniMaxM2ForCausalLM diff --git a/python/tokenspeed/runtime/models/qwen2.py b/python/tokenspeed/runtime/models/qwen2.py new file mode 100644 index 0000000..096d351 --- /dev/null +++ b/python/tokenspeed/runtime/models/qwen2.py @@ -0,0 +1,495 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Inference-only Qwen2 model compatible with HuggingFace weights.""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +import torch +from torch import nn + +from tokenspeed.runtime.configs.qwen2_config import Qwen2Config +from tokenspeed.runtime.configs.utils import get_rope_theta +from tokenspeed.runtime.distributed.comm_ops import all_reduce +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.layers.activation import SiluAndMul +from tokenspeed.runtime.layers.layernorm import RMSNorm +from tokenspeed.runtime.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from tokenspeed.runtime.layers.paged_attention import PagedAttention +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.layers.rotary_embedding import get_rope +from tokenspeed.runtime.layers.utils import get_layer_id +from tokenspeed.runtime.layers.vocab_parallel_embedding import VocabParallelEmbedding +from tokenspeed.runtime.model_loader.weight_utils import ( + default_weight_loader, + kv_cache_scales_loader, +) +from tokenspeed.runtime.models.base import BaseCausalLM +from tokenspeed.runtime.models.utils import validate_attention_partition +from tokenspeed.runtime.utils import add_prefix, make_layers +from tokenspeed.runtime.utils.env import global_server_args_dict + + +class Qwen2MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: QuantizationConfig | None = None, + tp_rank: int | None = None, + tp_size: int | None = None, + tp_group: tuple[int, ...] | None = None, + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + tp_rank=tp_rank, + tp_size=tp_size, + tp_group=tp_group, + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=False, + tp_rank=tp_rank, + tp_size=tp_size, + tp_group=tp_group, + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. " + "Only silu is supported for now." + ) + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class Qwen2Attention(nn.Module): + def __init__( + self, + config: Qwen2Config, + mapping: Mapping, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + layer_id: int = 0, + rope_theta: float = 1000000, + rope_scaling: dict[str, Any] | None = None, + head_dim: int | None = None, + max_position_embeddings: int = 32768, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.mapping = mapping + self.hidden_size = hidden_size + self.tp_rank = self.mapping.attn.tp_rank + self.tp_size = self.mapping.attn.tp_size + self.total_num_heads = num_heads + self.total_num_kv_heads = num_kv_heads + validate_attention_partition( + self.total_num_heads, + self.total_num_kv_heads, + self.tp_size, + ) + self.num_heads = self.total_num_heads // self.tp_size + self.num_kv_heads = max(1, self.total_num_kv_heads // self.tp_size) + self.head_dim = head_dim or hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + # Qwen2 uses biases on Q/K/V projections but not on the output projection. + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=True, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + reduce_results=False, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + ) + self.attn = PagedAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + cos_sin: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v, ctx, out_cache_loc) + if len(attn_output.size()) == 3: + attn_output = attn_output.reshape(attn_output.shape[0], -1) + output, _ = self.o_proj(attn_output) + return output + + +class Qwen2DecoderLayer(nn.Module): + def __init__( + self, + config: Qwen2Config, + mapping: Mapping, + layer_id: int = 0, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.mapping = mapping + if self.mapping.attn.tp_size != self.mapping.dense.tp_size: + raise ValueError( + "Qwen2 does not use CommManager and assumes attn_tp_size == dense_tp_size" + ) + self.hidden_size = config.hidden_size + rope_theta = get_rope_theta(config, 1000000) + rope_scaling = getattr(config, "rope_scaling", None) + max_position_embeddings = getattr(config, "max_position_embeddings", 32768) + head_dim = getattr(config, "head_dim", None) + self.self_attn = Qwen2Attention( + config=config, + mapping=self.mapping, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + layer_id=layer_id, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + head_dim=head_dim, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + ) + self.mlp = Qwen2MLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + tp_rank=self.mapping.dense.tp_rank, + tp_size=self.mapping.dense.tp_size, + tp_group=self.mapping.dense.tp_group, + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + residual: torch.Tensor | None, + cos_sin: tuple[torch.Tensor, torch.Tensor] | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + elif ( + ctx.input_num_tokens > global_server_args_dict["comm_fusion_max_num_tokens"] + ): + hidden_states = all_reduce(hidden_states, self.mapping.dense.tp_group) + hidden_states, residual = self.input_layernorm(hidden_states, residual) + else: + hidden_states, residual, *_ = ( + self.input_layernorm.forward_with_allreduce_fusion( + self.mapping.dense.tp_rank, + self.mapping.dense.tp_group, + hidden_states, + residual, + ) + ) + + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ctx=ctx, + out_cache_loc=out_cache_loc, + cos_sin=cos_sin, + ) + + if ctx.input_num_tokens > global_server_args_dict["comm_fusion_max_num_tokens"]: + hidden_states = all_reduce(hidden_states, self.mapping.attn.tp_group) + hidden_states, residual = self.post_attention_layernorm( + hidden_states, residual + ) + else: + hidden_states, residual, *_ = ( + self.post_attention_layernorm.forward_with_allreduce_fusion( + self.mapping.attn.tp_rank, + self.mapping.attn.tp_group, + hidden_states, + residual, + ) + ) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class Qwen2Model(nn.Module): + def __init__( + self, + config: Qwen2Config, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + decoder_layer_type: type[nn.Module] = None, + ) -> None: + super().__init__() + self.mapping = mapping + self.config = config + self.padding_idx = getattr(config, "pad_token_id", None) + self.vocab_size = config.vocab_size + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + decoder_layer_type = decoder_layer_type or Qwen2DecoderLayer + self.layers = make_layers( + config.num_hidden_layers, + lambda idx, prefix: decoder_layer_type( + config=config, + mapping=self.mapping, + layer_id=idx, + quant_config=quant_config, + ), + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + if hasattr(self.config, "scale_emb"): + return self.embed_tokens(input_ids) * self.config.scale_emb + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + input_embeds: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + residual = None + + for i in range(len(self.layers)): + layer = self.layers[i] + hidden_states, residual = layer( + positions, + hidden_states, + ctx, + out_cache_loc, + residual, + cos_sin=None, + ) + if ctx.input_num_tokens > global_server_args_dict["comm_fusion_max_num_tokens"]: + hidden_states = all_reduce(hidden_states, self.mapping.dense.tp_group) + hidden_states, _ = self.norm(hidden_states, residual) + else: + hidden_states, *_ = self.norm.forward_with_allreduce_fusion( + self.mapping.dense.tp_rank, + self.mapping.dense.tp_group, + hidden_states, + residual, + ) + return hidden_states, None + + def load_kv_cache_scales(self, quantization_param_path: str) -> None: + tp_size = self.mapping.attn.tp_size + tp_rank = self.mapping.attn.tp_rank + for layer_idx, scaling_factor in kv_cache_scales_loader( + quantization_param_path, + tp_rank, + tp_size, + self.config.num_hidden_layers, + self.config.__class__.model_type, + ): + if not isinstance(self.layers[layer_idx], nn.Identity): + layer_self_attn = self.layers[layer_idx].self_attn + if hasattr(layer_self_attn.attn, "k_scale"): + layer_self_attn.attn.k_scale = scaling_factor + layer_self_attn.attn.v_scale = scaling_factor + else: + raise RuntimeError( + "Self attention has no KV cache scaling " "factor attribute!" + ) + + +class Qwen2ForCausalLM(BaseCausalLM): + model_cls = Qwen2Model + + default_bitsandbytes_target_modules = [ + ".gate_proj.", + ".down_proj.", + ".up_proj.", + ".q_proj.", + ".k_proj.", + ".v_proj.", + ".o_proj.", + ] + bitsandbytes_stacked_params_mapping = { + "q_proj": ("qkv_proj", 0), + "k_proj": ("qkv_proj", 1), + "v_proj": ("qkv_proj", 2), + "gate_proj": ("gate_up_proj", 0), + "up_proj": ("gate_up_proj", 1), + } + + def __init__( + self, + config: Qwen2Config, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + ) -> None: + super().__init__( + config=config, + mapping=mapping, + quant_config=quant_config, + ) + + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.get_input_embeddings(input_ids) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + if "Embedding" in self.config.name_or_path: + name = add_prefix(name, "model") + layer_id = get_layer_id(name) + if ( + layer_id is not None + and hasattr(self.model, "start_layer") + and ( + layer_id < self.model.start_layer + or layer_id >= self.model.end_layer + ) + ): + continue + + if "rotary_emb.inv_freq" in name or "projector" in name: + continue + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + continue + if self.config.tie_word_embeddings and "lm_head.weight" in name: + continue + if name.startswith("model.vision_tower") and name not in params_dict: + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + def get_embed_and_head(self): + return self.model.embed_tokens.weight, self.lm_head.weight + + def set_embed_and_head(self, embed, head): + del self.model.embed_tokens.weight + del self.lm_head.weight + self.model.embed_tokens.weight = embed + self.lm_head.weight = head + torch.cuda.empty_cache() + torch.cuda.synchronize() + + def load_kv_cache_scales(self, quantization_param_path: str) -> None: + self.model.load_kv_cache_scales(quantization_param_path) + + +EntryClass = Qwen2ForCausalLM diff --git a/python/tokenspeed/runtime/models/qwen3.py b/python/tokenspeed/runtime/models/qwen3.py new file mode 100755 index 0000000..072b038 --- /dev/null +++ b/python/tokenspeed/runtime/models/qwen3.py @@ -0,0 +1,541 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Inference-only Qwen2 model compatible with HuggingFace weights.""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +import torch +from tokenspeed_kernel.ops.layernorm.triton import qk_rmsnorm +from torch import nn + +from tokenspeed.runtime.configs.qwen3_config import Qwen3Config +from tokenspeed.runtime.configs.utils import get_rope_theta +from tokenspeed.runtime.distributed.comm_ops import all_reduce +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.layers.activation import SiluAndMul +from tokenspeed.runtime.layers.layernorm import RMSNorm +from tokenspeed.runtime.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from tokenspeed.runtime.layers.paged_attention import PagedAttention +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.layers.rotary_embedding import get_rope +from tokenspeed.runtime.layers.utils import get_layer_id +from tokenspeed.runtime.layers.vocab_parallel_embedding import VocabParallelEmbedding +from tokenspeed.runtime.model_loader.weight_utils import ( + default_weight_loader, + kv_cache_scales_loader, +) +from tokenspeed.runtime.models.base import BaseCausalLM +from tokenspeed.runtime.models.utils import validate_attention_partition +from tokenspeed.runtime.utils import add_prefix, make_layers +from tokenspeed.runtime.utils.env import global_server_args_dict + + +class Qwen3MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: QuantizationConfig | None = None, + tp_rank: int | None = None, + tp_size: int | None = None, + tp_group: tuple[int, ...] | None = None, + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + tp_rank=tp_rank, + tp_size=tp_size, + tp_group=tp_group, + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=False, + tp_rank=tp_rank, + tp_size=tp_size, + tp_group=tp_group, + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. " + "Only silu is supported for now." + ) + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class Qwen3Attention(nn.Module): + def __init__( + self, + config: Qwen3Config, + mapping: Mapping, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + layer_id: int = 0, + rope_theta: float = 1000000, + rope_scaling: dict[str, Any] | None = None, + head_dim: int | None = None, + max_position_embeddings: int = 32768, + quant_config: QuantizationConfig | None = None, + rms_norm_eps: float = None, + attention_bias: bool = False, + prefix: str = "", + ) -> None: + super().__init__() + self.mapping = mapping + self.hidden_size = hidden_size + self.tp_rank = self.mapping.attn.tp_rank + self.tp_size = self.mapping.attn.tp_size + self.total_num_heads = num_heads + self.total_num_kv_heads = num_kv_heads + validate_attention_partition( + self.total_num_heads, + self.total_num_kv_heads, + self.tp_size, + ) + self.num_heads = self.total_num_heads // self.tp_size + self.num_kv_heads = max(1, self.total_num_kv_heads // self.tp_size) + self.head_dim = head_dim or hidden_size // self.total_num_heads + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps) + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=attention_bias, + quant_config=quant_config, + prefix=add_prefix("qkv_proj", prefix), + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + hidden_size, + bias=attention_bias, + quant_config=quant_config, + prefix=add_prefix("o_proj", prefix), + reduce_results=False, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.head_dim, + max_position=max_position_embeddings, + base=rope_theta, + rope_scaling=rope_scaling, + ) + self.attn = PagedAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + ) + + def _apply_qk_norm( + self, q: torch.Tensor, k: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + return qk_rmsnorm( + q, + k, + self.q_norm.weight.data, + self.k_norm.weight.data, + self.q_norm.variance_epsilon, + ) + + def _rotate_half(self, x): + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + def _apply_rotary_pos_emb(self, t, cos, sin): + return (t * cos) + self._rotate_half(t) * sin + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + cos_sin: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self._apply_qk_norm(q, k) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v, ctx, out_cache_loc) + if len(attn_output.size()) == 3: + attn_output = attn_output.reshape(attn_output.shape[0], -1) + output, _ = self.o_proj(attn_output) + return output + + +class Qwen3DecoderLayer(nn.Module): + def __init__( + self, + config: Qwen3Config, + mapping: Mapping, + layer_id: int = 0, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.mapping = mapping + if self.mapping.attn.tp_size != self.mapping.dense.tp_size: + raise ValueError( + "Qwen3 does not use CommManager and assumes attn_tp_size == dense_tp_size" + ) + self.hidden_size = config.hidden_size + rope_theta = get_rope_theta(config, 1000000) + rope_scaling = getattr(config, "rope_scaling", None) + max_position_embeddings = getattr(config, "max_position_embeddings", 32768) + head_dim = getattr(config, "head_dim", None) + self.self_attn = Qwen3Attention( + config=config, + mapping=self.mapping, + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + layer_id=layer_id, + rope_theta=rope_theta, + rope_scaling=rope_scaling, + head_dim=head_dim, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + rms_norm_eps=config.rms_norm_eps, + attention_bias=config.attention_bias, + prefix=add_prefix("self_attn", prefix), + ) + self.mlp = Qwen3MLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + tp_rank=self.mapping.dense.tp_rank, + tp_size=self.mapping.dense.tp_size, + tp_group=self.mapping.dense.tp_group, + ) + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + residual: torch.Tensor | None, + cos_sin: tuple[torch.Tensor, torch.Tensor] | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Self Attention + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + elif ( + ctx.input_num_tokens > global_server_args_dict["comm_fusion_max_num_tokens"] + ): + hidden_states = all_reduce(hidden_states, self.mapping.dense.tp_group) + hidden_states, residual = self.input_layernorm(hidden_states, residual) + else: + hidden_states, residual, *_ = ( + self.input_layernorm.forward_with_allreduce_fusion( + self.mapping.dense.tp_rank, + self.mapping.dense.tp_group, + hidden_states, + residual, + ) + ) + + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ctx=ctx, + out_cache_loc=out_cache_loc, + cos_sin=cos_sin, + ) + + # Fully Connected + if ctx.input_num_tokens > global_server_args_dict["comm_fusion_max_num_tokens"]: + hidden_states = all_reduce(hidden_states, self.mapping.attn.tp_group) + hidden_states, residual = self.post_attention_layernorm( + hidden_states, residual + ) + else: + hidden_states, residual, *_ = ( + self.post_attention_layernorm.forward_with_allreduce_fusion( + self.mapping.attn.tp_rank, + self.mapping.attn.tp_group, + hidden_states, + residual, + ) + ) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + +class Qwen3Model(nn.Module): + def __init__( + self, + config: Qwen3Config, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + decoder_layer_type: type[nn.Module] = None, + ) -> None: + super().__init__() + self.mapping = mapping + self.config = config + self.padding_idx = getattr(config, "pad_token_id", None) + self.vocab_size = config.vocab_size + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + decoder_layer_type = decoder_layer_type or Qwen3DecoderLayer + self.layers = make_layers( + config.num_hidden_layers, + lambda idx, prefix: decoder_layer_type( + config=config, + mapping=self.mapping, + layer_id=idx, + quant_config=quant_config, + ), + ) + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + if hasattr(self.config, "scale_emb"): + return self.embed_tokens(input_ids) * self.config.scale_emb + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + input_embeds: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + residual = None + + for i in range(len(self.layers)): + layer = self.layers[i] + hidden_states, residual = layer( + positions, + hidden_states, + ctx, + out_cache_loc, + residual, + cos_sin=None, + ) + if ctx.input_num_tokens > global_server_args_dict["comm_fusion_max_num_tokens"]: + hidden_states = all_reduce(hidden_states, self.mapping.dense.tp_group) + hidden_states, _ = self.norm(hidden_states, residual) + else: + hidden_states, *_ = self.norm.forward_with_allreduce_fusion( + self.mapping.dense.tp_rank, + self.mapping.dense.tp_group, + hidden_states, + residual, + ) + return hidden_states, None + + def load_kv_cache_scales(self, quantization_param_path: str) -> None: + tp_size = self.mapping.attn.tp_size + tp_rank = self.mapping.attn.tp_rank + for layer_idx, scaling_factor in kv_cache_scales_loader( + quantization_param_path, + tp_rank, + tp_size, + self.config.num_hidden_layers, + self.config.__class__.model_type, + ): + if not isinstance(self.layers[layer_idx], nn.Identity): + layer_self_attn = self.layers[layer_idx].self_attn + if hasattr(layer_self_attn.attn, "k_scale"): + layer_self_attn.attn.k_scale = scaling_factor + layer_self_attn.attn.v_scale = scaling_factor + else: + raise RuntimeError( + "Self attention has no KV cache scaling " "factor attribute!" + ) + + +class Qwen3ForCausalLM(BaseCausalLM): + model_cls = Qwen3Model + + # BitandBytes specific attributes + default_bitsandbytes_target_modules = [ + ".gate_proj.", + ".down_proj.", + ".up_proj.", + ".q_proj.", + ".k_proj.", + ".v_proj.", + ".o_proj.", + ] + bitsandbytes_stacked_params_mapping = { + # shard_name, weight_name, index + "q_proj": ("qkv_proj", 0), + "k_proj": ("qkv_proj", 1), + "v_proj": ("qkv_proj", 2), + "gate_proj": ("gate_up_proj", 0), + "up_proj": ("gate_up_proj", 1), + } + + def __init__( + self, + config: Qwen3Config, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + ) -> None: + super().__init__( + config=config, + mapping=mapping, + quant_config=quant_config, + ) + + def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.get_input_embeddings(input_ids) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + """ + 'model.layers.0.self_attn.q_norm.weight', + 'model.layers.0.self_attn.k_norm.weight', + 'model.layers.0.self_attn.qkv_proj.weight', + 'model.layers.0.self_attn.o_proj.weight', + 'model.layers.0.mlp.gate_up_proj.weight', + 'model.layers.0.mlp.down_proj.weight', + 'model.layers.0.input_layernorm.weight', + 'model.layers.0.post_attention_layernorm.weight' + """ + params_dict = dict(self.named_parameters()) + for name, loaded_weight in weights: + if "Embedding" in self.config.name_or_path: + name = add_prefix(name, "model") + layer_id = get_layer_id(name) + if ( + layer_id is not None + and hasattr(self.model, "start_layer") + and ( + layer_id < self.model.start_layer + or layer_id >= self.model.end_layer + ) + ): + continue + + if "rotary_emb.inv_freq" in name or "projector" in name: + continue + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + # Models trained using ColossalAI may include these tensors in + # the checkpoint. Skip them. + continue + if self.config.tie_word_embeddings and "lm_head.weight" in name: + continue + if name.startswith("model.vision_tower") and name not in params_dict: + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + def get_embed_and_head(self): + return self.model.embed_tokens.weight, self.lm_head.weight + + def set_embed_and_head(self, embed, head): + del self.model.embed_tokens.weight + del self.lm_head.weight + self.model.embed_tokens.weight = embed + self.lm_head.weight = head + torch.cuda.empty_cache() + torch.cuda.synchronize() + + def load_kv_cache_scales(self, quantization_param_path: str) -> None: + self.model.load_kv_cache_scales(quantization_param_path) + + +EntryClass = Qwen3ForCausalLM diff --git a/python/tokenspeed/runtime/models/qwen3_5.py b/python/tokenspeed/runtime/models/qwen3_5.py new file mode 100644 index 0000000..bc70aa1 --- /dev/null +++ b/python/tokenspeed/runtime/models/qwen3_5.py @@ -0,0 +1,1772 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Inference-only Qwen3.5 model and Qwen3.5 MoE model compatible with HuggingFace weights.""" + +from __future__ import annotations + +import logging +from collections.abc import Iterable + +import torch +import torch.nn as nn +import triton +import triton.language as tl +from tokenspeed_kernel.ops.activation.triton import sigmoid_mul +from tokenspeed_kernel.ops.layernorm.triton import ( + fused_qk_rmsnorm_rope_gate, + qk_rmsnorm, +) + +# Configs +from tokenspeed.runtime.configs.paged_cache_spec import FULL_ATTENTION +from tokenspeed.runtime.configs.qwen3_5_config import ( + Qwen3_5Config, + Qwen3_5TextConfig, +) +from tokenspeed.runtime.configs.utils import get_rope_parameters + +# Distributed +from tokenspeed.runtime.distributed.comm_manager import CommManager +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.execution.context import ForwardContext + +# Layers - Attention +from tokenspeed.runtime.layers.attention.linear.layernorm_gated import ( + RMSNorm as RMSNormGated, +) + +# Layers - Others +from tokenspeed.runtime.layers.layernorm import GemmaRMSNorm + +# Layers - Linear +from tokenspeed.runtime.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from tokenspeed.runtime.layers.logits_processor import LogitsMetadata +from tokenspeed.runtime.layers.moe import ( + ExpertCheckpointSchema, + build_moe_checkpoint_loader, +) +from tokenspeed.runtime.layers.paged_attention import PagedAttention +from tokenspeed.runtime.layers.parameter import ( + BlockQuantScaleParameter, + PerTensorScaleParameter, +) +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.layers.rotary_embedding import get_rope +from tokenspeed.runtime.layers.vocab_parallel_embedding import VocabParallelEmbedding +from tokenspeed.runtime.model_loader.weight_utils import ( + default_weight_loader, + mamba_v2_sharded_weight_loader, + sharded_weight_loader, +) +from tokenspeed.runtime.models.base import BaseCausalLM +from tokenspeed.runtime.models.qwen3_5_moe import ( + Qwen3_5MoeMLP, + Qwen3_5MoeSparseMoeBlock, +) +from tokenspeed.runtime.models.qwen3_vision import Qwen3VLMoeVisionModel +from tokenspeed.runtime.models.utils import validate_attention_partition +from tokenspeed.runtime.moe.distribution_recorder import ( + get_global_expert_distribution_recorder, +) +from tokenspeed.runtime.moe.expert_location import ModelConfigForExpertLocation +from tokenspeed.runtime.multimodal.embedder import ( + EncoderSpec, + VisionEmbedder, + pad_input_tokens, +) +from tokenspeed.runtime.multimodal.encoder_cudagraph import ( + EncoderCudaGraphWrapper, + VisionEncoderCudaGraphAdapter, +) +from tokenspeed.runtime.multimodal.inputs import ( + Modality, + MultimodalDataItem, + MultimodalInputs, +) +from tokenspeed.runtime.utils import ( + add_prefix, + make_layers, + set_weight_attrs, +) +from tokenspeed.runtime.utils.env import envs + +logger = logging.getLogger(__name__) + + +class Qwen3_5GatedDeltaNet(nn.Module): + def __init__( + self, + config: Qwen3_5TextConfig, + mapping: Mapping, + layer_id: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.mapping = mapping + self.attn_tp_rank = mapping.attn.tp_rank + self.attn_tp_size = mapping.attn.tp_size + self.attn_tp_group = mapping.attn.tp_group + self.hidden_size = config.hidden_size + self.num_v_heads = config.linear_num_value_heads + self.num_k_heads = config.linear_num_key_heads + self.head_k_dim = config.linear_key_head_dim + self.head_v_dim = config.linear_value_head_dim + self.key_dim = self.head_k_dim * self.num_k_heads + self.value_dim = self.head_v_dim * self.num_v_heads + + self.conv_kernel_size = config.linear_conv_kernel_dim + self.layer_id = layer_id + self.activation = config.hidden_act + self.layer_norm_epsilon = config.rms_norm_eps + + # Conv1d layer + self.conv_dim = self.key_dim * 2 + self.value_dim + self.conv1d = ColumnParallelLinear( + input_size=self.conv_kernel_size, + output_size=self.conv_dim, + bias=False, + quant_config=None, + tp_rank=self.attn_tp_rank, + tp_size=self.attn_tp_size, + tp_group=self.attn_tp_group, + prefix=add_prefix("conv1d", prefix), + ) + self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1) + + self.in_proj_qkvzba = MergedColumnParallelLinear( + input_size=self.hidden_size, + output_sizes=[ + self.key_dim, + self.key_dim, + self.value_dim, + self.value_dim, + self.num_v_heads, + self.num_v_heads, + ], + bias=False, + quant_config=quant_config, + tp_rank=self.attn_tp_rank, + tp_size=self.attn_tp_size, + tp_group=self.attn_tp_group, + prefix=add_prefix("in_proj_qkvzba", prefix), + ) + self._qkvz_dim = (self.key_dim * 2 + self.value_dim * 2) // self.attn_tp_size + self._ba_dim = (self.num_v_heads * 2) // self.attn_tp_size + + # Override weight loaders for packed checkpoint format. + # Important: for FP8, this must cover not only `.weight` but also + # `weight_scale_inv` / `weight_scale` / `input_scale` if present. + self._bind_packed_weight_loaders(self.in_proj_qkvzba) + + # Conv1d weight loader setup + query_key_settings = (self.key_dim, 0, False) + value_settings = (self.value_dim, 0, False) + + delattr(self.conv1d.weight, "weight_loader") + set_weight_attrs( + self.conv1d.weight, + { + "weight_loader": mamba_v2_sharded_weight_loader( + [ + query_key_settings, + query_key_settings, + value_settings, + ], + self.attn_tp_size, + self.attn_tp_rank, + ) + }, + ) + + # State parameters + self.dt_bias = nn.Parameter( + torch.ones(self.num_v_heads // self.attn_tp_size), + ) + self.A_log = nn.Parameter( + torch.empty(self.num_v_heads // self.attn_tp_size), + ) + + set_weight_attrs( + self.A_log, {"weight_loader": sharded_weight_loader(0, self.attn_tp_rank)} + ) + set_weight_attrs( + self.dt_bias, {"weight_loader": sharded_weight_loader(0, self.attn_tp_rank)} + ) + + conv_weights = self.conv1d.weight.view( + self.conv1d.weight.size(0), self.conv1d.weight.size(2) + ) + + self.conv_weights = conv_weights + + # Normalization layer + self.norm = RMSNormGated( + self.head_v_dim, + eps=self.layer_norm_epsilon, + group_size=None, + norm_before_gate=True, + device=torch.get_device_module().current_device(), + dtype=config.dtype, + ) + + # Output projection + self.out_proj = RowParallelLinear( + self.value_dim, + self.hidden_size, + bias=False, + input_is_parallel=True, + reduce_results=False, + quant_config=quant_config, + tp_rank=self.attn_tp_rank, + tp_size=self.attn_tp_size, + tp_group=self.attn_tp_group, + prefix=add_prefix("out_proj", prefix), + ) + + @staticmethod + def _override_weight_loader(param, loader): + """Robustly override loader for: + 1) BaseWeightParameter subclasses: real storage is `_weight_loader` + 2) regular Parameters that already have mutable `weight_loader` + 3) regular Parameters without `weight_loader` yet + """ + if hasattr(param, "_weight_loader"): + # FP8 / quantized BaseWeightParameter path + param._weight_loader = loader + return + + if hasattr(param, "weight_loader"): + # Regular parameter/tensor that already has a mutable attr. + # Do NOT call set_weight_attrs here; overwriting an existing + # attribute is rejected. + param.weight_loader = loader + return + + # Fresh attribute on a normal tensor/Parameter + set_weight_attrs(param, {"weight_loader": loader}) + + def _bind_packed_weight_loaders(self, module): + """Bind packed-checkpoint-aware loaders to all relevant params of a merged module.""" + for attr_name in ("weight", "weight_scale_inv", "weight_scale", "input_scale"): + param = getattr(module, attr_name, None) + if param is None: + continue + original_loader = getattr(param, "weight_loader", None) + if original_loader is None: + continue + wrapped_loader = self._make_packed_weight_loader(module, original_loader) + self._override_weight_loader(param, wrapped_loader) + + @staticmethod + def _get_split_sizes_for_param(module, param, loaded_shard_id): + """Return checkpoint-side split sizes for this param type.""" + if isinstance(param, BlockQuantScaleParameter): + # Split by output blocks, not raw output sizes. + block_n, _ = module.quant_method.quant_config.weight_block_size + block_n = 1 if getattr(param, "format_ue8m0", False) else block_n + return [ + (module.output_sizes[idx] + block_n - 1) // block_n + for idx in loaded_shard_id + ] + + if isinstance(param, PerTensorScaleParameter): + # One logical scale per logical shard. + return [1 for _ in loaded_shard_id] + + # Normal weight / non-block quant tensor + return [module.output_sizes[idx] for idx in loaded_shard_id] + + @classmethod + def _make_packed_weight_loader(cls, module, original_weight_loader): + """Wrap the param's original loader so split checkpoints: + - in_proj_qkv + in_proj_z + in_proj_b + in_proj_a -> merged in_proj_qkvzba + can load correctly for both normal and FP8 params. + """ + + def weight_loader(param, loaded_weight, loaded_shard_id=None): + # Only intercept split-checkpoint tuple shards. + # int shard_id and None should preserve original behavior. + if isinstance(loaded_shard_id, tuple): + split_sizes = cls._get_split_sizes_for_param( + module, param, loaded_shard_id + ) + + if len(loaded_weight.shape) == 0: + # Scalar only makes sense for a single logical shard. + if len(split_sizes) != 1 or split_sizes[0] != 1: + raise ValueError( + f"Unexpected scalar for tuple shard load: " + f"{loaded_shard_id=}, {split_sizes=}" + ) + chunks = [loaded_weight.reshape(1)] + else: + split_dim = getattr(param, "output_dim", 0) + chunks = loaded_weight.split(split_sizes, dim=split_dim) + + if len(chunks) != len(loaded_shard_id): + raise ValueError( + f"Chunk/shard mismatch: {len(chunks)=}, " + f"{len(loaded_shard_id)=}, {split_sizes=}" + ) + + for idx, chunk in zip(loaded_shard_id, chunks): + # Delegate each chunk to the param's original int-shard loader. + original_weight_loader(param, chunk, idx) + return + + return original_weight_loader(param, loaded_weight, loaded_shard_id) + + return weight_loader + + def fix_query_key_value_ordering( + self, + mixed_qkvz: torch.Tensor, + mixed_ba: torch.Tensor, + ): + """ + Derives `query`, `key` and `value` tensors from `mixed_qkvzba`. + """ + k_tp = self.key_dim // self.attn_tp_size + v_tp = self.value_dim // self.attn_tp_size + nv_tp = self.num_v_heads // self.attn_tp_size + + # Directly split, no head group reshape + query, key, value, z = mixed_qkvz.split([k_tp, k_tp, v_tp, v_tp], dim=-1) + b, a = mixed_ba.split([nv_tp, nv_tp], dim=-1) + + # value / z reshape to (seq, num_v_heads/tp, head_v_dim) + value = value.reshape(value.size(0), -1, self.head_v_dim) + z = z.reshape(z.size(0), -1, self.head_v_dim) + + return query, key, value, z, b, a + + def _forward_input_proj(self, hidden_states: torch.Tensor): + projected_all, _ = self.in_proj_qkvzba(hidden_states) + projected_states_qkvz, projected_states_ba = projected_all.split( + [self._qkvz_dim, self._ba_dim], dim=-1 + ) + return projected_states_qkvz, projected_states_ba + + def forward( + self, + hidden_states: torch.Tensor, + ctx: ForwardContext, + ): + seq_len, _ = hidden_states.shape + + projected_states_qkvz, projected_states_ba = self._forward_input_proj( + hidden_states + ) + + if self.num_v_heads // self.num_k_heads in [1, 2, 4]: + mixed_qkv, z, b, a = fused_qkvzba_split_reshape_cat_contiguous( + projected_states_qkvz, + projected_states_ba, + triton.cdiv(self.num_k_heads, self.attn_tp_size), + triton.cdiv(self.num_v_heads, self.attn_tp_size), + self.head_k_dim, + self.head_v_dim, + ) + else: + query, key, value, z, b, a = self.fix_query_key_value_ordering( + projected_states_qkvz, projected_states_ba + ) + query, key, value = map( + lambda x: x.reshape(x.shape[0], -1), (query, key, value) + ) + mixed_qkv = torch.cat((query, key, value), dim=-1) + + kwargs = { + "mixed_qkv": mixed_qkv, + "conv_weights": self.conv_weights, + "bias": self.conv1d.bias, + "activation": self.activation, + "key_dim": self.key_dim, + "value_dim": self.value_dim, + "attention_tp_size": self.attn_tp_size, + "head_k_dim": self.head_k_dim, + "head_v_dim": self.head_v_dim, + "a": a, + "b": b, + "A_log": self.A_log, + "dt_bias": self.dt_bias, + "layer_id": self.layer_id, + "seq_len": seq_len, + "z": z, + } + + core_attn_out = ctx.attn_backend.forward( + q=None, + k=None, + v=None, + layer=None, + out_cache_loc=None, + token_to_kv_pool=ctx.token_to_kv_pool, + forward_mode=ctx.forward_mode, + bs=ctx.bs, + **kwargs, + ) + + z_shape_og = z.shape + core_attn_out = core_attn_out.reshape(-1, core_attn_out.shape[-1]) + z = z.reshape(-1, z.shape[-1]) + core_attn_out = self.norm(core_attn_out, z) + core_attn_out = core_attn_out.reshape(z_shape_og) + core_attn_out = core_attn_out.reshape(*core_attn_out.shape[:-2], -1) + output, _ = self.out_proj(core_attn_out) + return output + + +class Qwen3_5LinearDecoderLayer(nn.Module): + """Qwen3.5 Decoder Layer with Linear Attention (GatedDeltaNet).""" + + def __init__( + self, + config: Qwen3_5TextConfig, + mapping: Mapping, + layer_id: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + alt_stream: torch.cuda.Stream | None = None, + ) -> None: + super().__init__() + self.config = config + self.mapping = mapping + self.layer_id = layer_id + + linear_attn_quant_config = ( + None + if quant_config and quant_config.get_name() in ("fp8", "nvfp4") + else quant_config + ) + self.linear_attn = Qwen3_5GatedDeltaNet( + config, mapping, layer_id, linear_attn_quant_config, prefix=prefix + ) + + # Determine the MLP type based on the model type + # Qwen3.5 use all layers for MLP / Qwen3.5-MoE use sparse MoE blocks + if config.model_type == "qwen3_5_moe_text": + self.mlp = Qwen3_5MoeSparseMoeBlock( + config=config, + mapping=self.mapping, + quant_config=quant_config, + layer_index=layer_id, + alt_stream=alt_stream, + prefix=add_prefix("mlp", prefix.replace(".linear_attn", "")), + ) + is_moe = True + elif config.model_type == "qwen3_5_text": + self.mlp = Qwen3_5MoeMLP( + mapping=self.mapping, + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=add_prefix("mlp", prefix.replace(".linear_attn", "")), + ) + is_moe = False + else: + raise ValueError(f"Invalid model type: {config.model_type}") + + self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = GemmaRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + self.is_moe = is_moe + self.comm_manager = CommManager( + mapping=self.mapping, + layer_id=self.layer_id, + is_moe=is_moe, + prev_is_moe=is_moe, + input_layernorm=self.input_layernorm, + post_attn_layernorm=self.post_attention_layernorm, + ) + + def forward( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ctx: ForwardContext, + **kwargs, + ): + num_global_tokens, max_num_tokens_per_gpu = self.comm_manager.get_num_tokens( + ctx + ) + + if not ctx.forward_mode.is_idle(): + hidden_states, residual = self.comm_manager.input_reduce_norm( + hidden_states, residual + ) + hidden_states = self.comm_manager.pre_attn_comm(hidden_states, ctx) + + hidden_states = self.linear_attn( + hidden_states, + ctx, + ) + + hidden_states, residual = self.comm_manager.post_attn_reduce_norm( + hidden_states, residual, ctx + ) + + hidden_states = self.forward_mlp( + hidden_states, + residual, + ctx, + num_global_tokens, + max_num_tokens_per_gpu, + ) + + return hidden_states, residual + + def forward_mlp( + self, + hidden_states, + residual, + ctx: ForwardContext, + num_global_tokens, + max_num_tokens_per_gpu, + ): + if isinstance(self.mlp, Qwen3_5MoeSparseMoeBlock): + hidden_states = self.mlp( + hidden_states, num_global_tokens, max_num_tokens_per_gpu, ctx + ) + else: + hidden_states = self.comm_manager.pre_mlp_comm(hidden_states, ctx) + hidden_states = self.mlp(hidden_states) + hidden_states, residual = self.comm_manager.post_mlp_fused( + hidden_states, residual, ctx + ) + return hidden_states + + +class Qwen3_5AttentionDecoderLayer(nn.Module): + """Qwen3.5 Decoder Layer with Full Attention.""" + + def __init__( + self, + config: Qwen3_5TextConfig, + mapping: Mapping, + layer_id: int, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + alt_stream: torch.cuda.Stream | None = None, + ) -> None: + super().__init__() + self.config = config + self.mapping = mapping + self.hidden_size = config.hidden_size + self.attn_tp_rank = mapping.attn.tp_rank + self.attn_tp_size = mapping.attn.tp_size + self.attn_tp_group = mapping.attn.tp_group + self.total_num_heads = config.num_attention_heads + self.total_num_kv_heads = config.num_key_value_heads + validate_attention_partition( + self.total_num_heads, + self.total_num_kv_heads, + self.attn_tp_size, + ) + self.num_heads = self.total_num_heads // self.attn_tp_size + self.num_kv_heads = max(1, self.total_num_kv_heads // self.attn_tp_size) + self.head_dim = config.head_dim or (self.hidden_size // self.num_heads) + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + self.max_position_embeddings = getattr(config, "max_position_embeddings", 8192) + + self.rope_scaling = get_rope_parameters(config) + + self.rope_theta = self.rope_scaling.get("rope_theta", 10000) + self.partial_rotary_factor = self.rope_scaling.get("partial_rotary_factor", 1.0) + self.layer_id = layer_id + + self.attn_output_gate = getattr(config, "attn_output_gate", True) + if self.attn_output_gate: + logger.warning_once("using attn output gate!") + + self.rotary_emb = get_rope( + head_size=self.head_dim, + rotary_dim=self.head_dim, + max_position=self.max_position_embeddings, + rope_scaling=self.rope_scaling, + base=self.rope_theta, + partial_rotary_factor=self.partial_rotary_factor, + is_neox_style=True, + dtype=torch.get_default_dtype(), + ) + + attn_quant_config = ( + None + if quant_config and quant_config.get_name() == "nvfp4" + else quant_config + ) + + self.qkv_proj = QKVParallelLinear( + config.hidden_size, + self.head_dim, + self.total_num_heads * (1 + self.attn_output_gate), + self.total_num_kv_heads, + bias=False, + quant_config=attn_quant_config, + tp_rank=self.attn_tp_rank, + tp_size=self.attn_tp_size, + tp_group=self.attn_tp_group, + prefix=add_prefix("qkv_proj", prefix), + ) + + self.o_proj = RowParallelLinear( + self.total_num_heads * self.head_dim, + config.hidden_size, + bias=False, + quant_config=attn_quant_config, + reduce_results=False, + tp_rank=self.attn_tp_rank, + tp_size=self.attn_tp_size, + tp_group=self.attn_tp_group, + prefix=add_prefix("o_proj", prefix), + ) + + self.attn = PagedAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + group_id=FULL_ATTENTION, + ) + + # Dense MLP for non-MoE variant + if config.model_type == "qwen3_5_text": + self.mlp = Qwen3_5MoeMLP( + mapping=self.mapping, + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=add_prefix("mlp", prefix.replace(".self_attn", "")), + ) + is_moe = False + elif config.model_type == "qwen3_5_moe_text": + self.mlp = Qwen3_5MoeSparseMoeBlock( + config=config, + mapping=self.mapping, + quant_config=quant_config, + layer_index=layer_id, + alt_stream=alt_stream, + prefix=add_prefix("mlp", prefix.replace(".self_attn", "")), + ) + is_moe = True + else: + raise ValueError(f"Invalid model type: {config.model_type}") + + self.input_layernorm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = GemmaRMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + self.q_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = GemmaRMSNorm(self.head_dim, eps=config.rms_norm_eps) + + self.is_moe = is_moe + self.comm_manager = CommManager( + mapping=self.mapping, + layer_id=self.layer_id, + is_moe=is_moe, + prev_is_moe=is_moe, + input_layernorm=self.input_layernorm, + post_attn_layernorm=self.post_attention_layernorm, + ) + + def _apply_qk_norm( + self, q: torch.Tensor, k: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + # qk_rmsnorm expects GemmaRMSNorm's effective gamma. + return qk_rmsnorm( + q, + k, + self.q_norm.gemma_weight, + self.k_norm.gemma_weight, + self.q_norm.variance_epsilon, + ) + + def _project_qkv_rope( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: + """qkv_proj + split + rope (+ optional gate). ``gate`` is ``None`` when ``attn_output_gate=False``.""" + qkv, _ = self.qkv_proj(hidden_states) + if self.attn_output_gate: + q_gate, k, v = qkv.split( + [self.q_size * 2, self.kv_size, self.kv_size], dim=-1 + ) + q, k, gate = fused_qk_rmsnorm_rope_gate( + q_gate, + k, + self.q_norm.gemma_weight, + self.k_norm.gemma_weight, + self.rotary_emb.cos_sin_cache, + positions, + self.q_norm.variance_epsilon, + self.num_heads, + self.num_kv_heads, + self.head_dim, + self.rotary_emb.rotary_dim, + ) + return q, k, v, gate + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + q, k = self._apply_qk_norm(q, k) + q, k = self.rotary_emb(positions, q, k) + return q, k, v, None + + def _attn( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gate: torch.Tensor | None, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + ) -> torch.Tensor: + """Backend attention call + optional gate apply. Subclasses override.""" + attn_output = self.attn(q, k, v, ctx, out_cache_loc) + if gate is not None: + sigmoid_mul(attn_output, gate) + return attn_output + + def self_attention( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + ) -> torch.Tensor: + """Full attention forward pass.""" + q, k, v, gate = self._project_qkv_rope(positions, hidden_states) + attn_output = self._attn(q, k, v, gate, ctx, out_cache_loc) + output, _ = self.o_proj(attn_output) + return output + + def _maybe_narrow_residual( + self, + residual: torch.Tensor, + ctx: ForwardContext, + ) -> torch.Tensor: + """Hook: subclasses narrow residual to match a sliced attn output.""" + return residual + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + **kwargs, + ): + num_global_tokens, max_num_tokens_per_gpu = self.comm_manager.get_num_tokens( + ctx + ) + + if not ctx.forward_mode.is_idle(): + hidden_states, residual = self.comm_manager.input_reduce_norm( + hidden_states, residual + ) + hidden_states = self.comm_manager.pre_attn_comm(hidden_states, ctx) + hidden_states = self.self_attention( + positions=positions, + hidden_states=hidden_states, + ctx=ctx, + out_cache_loc=out_cache_loc, + ) + residual = self._maybe_narrow_residual(residual, ctx) + hidden_states, residual = self.comm_manager.post_attn_reduce_norm( + hidden_states, residual, ctx + ) + + hidden_states = self.forward_mlp( + hidden_states, + residual, + ctx, + num_global_tokens, + max_num_tokens_per_gpu, + ) + + return hidden_states, residual + + def forward_mlp( + self, + hidden_states, + residual, + ctx: ForwardContext, + num_global_tokens, + max_num_tokens_per_gpu, + ): + if isinstance(self.mlp, Qwen3_5MoeSparseMoeBlock): + hidden_states = self.mlp( + hidden_states, num_global_tokens, max_num_tokens_per_gpu, ctx + ) + else: + hidden_states = self.comm_manager.pre_mlp_comm(hidden_states, ctx) + hidden_states = self.mlp(hidden_states) + hidden_states, residual = self.comm_manager.post_mlp_fused( + hidden_states, residual, ctx + ) + return hidden_states + + +class Qwen3_5ForCausalLM(nn.Module): + """Qwen3.5 Model with support for dense variant.""" + + ATTENTION_LAYER_CLS: type = Qwen3_5AttentionDecoderLayer + LINEAR_LAYER_CLS: type = Qwen3_5LinearDecoderLayer + + def __init__( + self, + config: Qwen3_5TextConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + self.mapping = mapping + self.hidden_size = config.hidden_size + + alt_stream = torch.cuda.Stream() + + # Embedding layer + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + org_num_embeddings=config.vocab_size, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + + layer_cls_by_type = { + "attention": self.ATTENTION_LAYER_CLS, + "linear_attention": self.LINEAR_LAYER_CLS, + } + + def get_layer(idx: int, prefix: str): + layer_type = config.layers_block_type[idx] + layer_class = layer_cls_by_type[layer_type] + if layer_type == "attention": + prefix = add_prefix("self_attn", prefix) + else: + prefix = add_prefix("linear_attn", prefix) + return layer_class( + config=config, + mapping=self.mapping, + layer_id=idx, + quant_config=quant_config, + prefix=prefix, + alt_stream=alt_stream, + ) + + self.layers = make_layers( + config.num_hidden_layers, + get_layer, + prefix=f"{prefix}.layers", + ) + + # Final normalization + self.norm = GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def get_input_embeddings(self) -> nn.Embedding: + return self.embed_tokens + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + input_embeds: torch.Tensor | None = None, + pp_proxy_tensors=None, + input_deepstack_embeds: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + # Initialize hidden states + if input_embeds is None: + # Only skip embedding allreduce when the first layer's fused + # allreduce+residual+norm will handle it + if self.layers[0].comm_manager.should_fuse(input_ids.shape[0]): + hidden_states = self.embed_tokens(input_ids, reduce_results=False) + residual = torch.zeros_like(hidden_states) + else: + hidden_states = self.embed_tokens(input_ids) + residual = None + else: + hidden_states = input_embeds + residual = None + + # Pass through decoder layers + for layer_idx in range(len(self.layers)): + layer = self.layers[layer_idx] + with get_global_expert_distribution_recorder().with_current_layer( + layer_idx + ): + hidden_states, residual = layer( + positions=positions, + hidden_states=hidden_states, + residual=residual, + ctx=ctx, + out_cache_loc=out_cache_loc, + ) + + # Process deepstack embeddings if provided + if ( + input_deepstack_embeds is not None + and input_deepstack_embeds.numel() > 0 + and layer_idx < 3 + ): + sep = self.hidden_size * layer_idx + hidden_states.add_( + input_deepstack_embeds[:, sep : sep + self.hidden_size] + ) + + # Apply final normalization with optional allreduce fusion + hidden_states, _ = layer.comm_manager.final_norm( + hidden_states, residual, ctx, self.norm + ) + + return hidden_states, None + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + # GDN (GatedDeltaNet) linear attention projections + # Split checkpoint format (separate qkv/z/b/a files) + ("in_proj_qkvzba.", "in_proj_qkv.", (0, 1, 2)), + ("in_proj_qkvzba.", "in_proj_z.", 3), + ("in_proj_qkvzba.", "in_proj_b.", 4), + ("in_proj_qkvzba.", "in_proj_a.", 5), + # Pre-packed checkpoint format (already merged qkvz and ba) + ("in_proj_qkvzba.", "in_proj_qkvz.", (0, 1, 2, 3)), + ("in_proj_qkvzba.", "in_proj_ba.", (4, 5)), + ] + + loaded_params: set[str] = set() + params_dict = dict(self.named_parameters(remove_duplicate=False)) + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + if "mtp" in name: + continue + if "visual" in name: + continue + if "language_model" in name: + name = name.replace(r"model.language_model.", r"model.") + if ".self_attn." in name: + name = name.replace(".self_attn", "") + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + + if "mlp.experts" in name: + continue + + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader") + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if name not in params_dict: + logger.warning("Parameter %s not found in params_dict", name) + continue + param = params_dict[name] + + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params + + +class Qwen3_5MoeForCausalLM(Qwen3_5ForCausalLM): + def __init__( + self, + config: Qwen3_5TextConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__( + config=config, mapping=mapping, quant_config=quant_config, prefix=prefix + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + # GDN (GatedDeltaNet) linear attention projections + # Split checkpoint format (separate qkv/z/b/a files) + ("in_proj_qkvzba.", "in_proj_qkv.", (0, 1, 2)), + ("in_proj_qkvzba.", "in_proj_z.", 3), + ("in_proj_qkvzba.", "in_proj_b.", 4), + ("in_proj_qkvzba.", "in_proj_a.", 5), + # Pre-packed checkpoint format (already merged qkvz and ba) + ("in_proj_qkvzba.", "in_proj_qkvz.", (0, 1, 2, 3)), + ("in_proj_qkvzba.", "in_proj_ba.", (4, 5)), + ] + + # Skip loading extra parameters for GPTQ/nvfp4 models. + ignore_suffixes = ( + ".bias", + "_bias", + ".k_scale", + "_k_scale", + ".v_scale", + "_v_scale", + ".weight_scale", + "_weight_scale", + ".input_scale", + "_input_scale", + ) + loaded_params: set[str] = set() + params_dict = dict(self.named_parameters(remove_duplicate=False)) + # MoE expert weights, scales, and activation scales are handled + # by the checkpoint loader. + moe_loader = build_moe_checkpoint_loader( + params_dict=params_dict, + expert_schema=ExpertCheckpointSchema( + gate_proj_name="gate_proj", + down_proj_name="down_proj", + up_proj_name="up_proj", + ), + fused_schema=ExpertCheckpointSchema( + gate_up_fused_name="gate_up_proj", + down_proj_name="down_proj", + ), + num_experts=self.config.num_experts, + ep_rank=self.mapping.moe.ep_rank, + ep_size=self.mapping.moe.ep_size, + ) + + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + if "mtp" in name: + continue + if "visual" in name: + continue + if "language_model" in name: + name = name.replace(r"model.language_model.", r"model.") + if ".self_attn." in name: + name = name.replace(".self_attn", "") + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + + if "mlp.experts" in name: + continue + name = name.replace(weight_name, param_name) + + # Skip loading extra parameters for GPTQ/nvfp4 models. + if name.endswith(ignore_suffixes) and name not in params_dict: + continue + + if name not in params_dict: + continue + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith((".bias", "_bias")) and name not in params_dict: + continue + if moe_loader.matches(name): + mapped_name = moe_loader.load(name, loaded_weight) + loaded_params.add(mapped_name) + continue + if moe_loader.is_expert_checkpoint_weight(name): + continue + + # Skip loading extra parameters for GPTQ/nvfp4 models. + if name.endswith(ignore_suffixes) and name not in params_dict: + continue + + if name in params_dict.keys(): + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + else: + logger.warning("Parameter %s not found in params_dict", name) + loaded_params.add(name) + + return loaded_params + + +class Qwen3_5ForConditionalGeneration(BaseCausalLM): + model_cls = Qwen3_5ForCausalLM + + def __init__( + self, + config: Qwen3_5Config, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + is_multimodal_active: bool = True, + mm_attention_backend: str | None = None, + ): + super().__init__( + config=config.text_config, + mapping=mapping, + quant_config=quant_config, + prefix=prefix, + encoder_only=getattr(config, "encoder_only", False), + ) + + rope_config = get_rope_parameters(self.config) + self.is_mrope_enabled = "mrope_section" in rope_config + self.is_multimodal_active = is_multimodal_active + if not self.is_multimodal_active: + self.visual = None + self.deepstack_visual_indexes = [] + self.num_deepstack_embeddings = 0 + self.vision_embedder = None + self.image_encoder = None + self.video_encoder = None + else: + self.visual = Qwen3VLMoeVisionModel( + config.vision_config, + quant_config=None, + norm_eps=getattr(config, "rms_norm_eps", 1e-6), + prefix=add_prefix("model.visual", prefix), + mapping=mapping, + mm_attention_backend=mm_attention_backend, + ) + self.deepstack_visual_indexes = self.visual.deepstack_visual_indexes + self.num_deepstack_embeddings = len(self.deepstack_visual_indexes) + # Encoder callables may be swapped to cudagraph wrappers by + # ModelExecutor. + self.vision_embedder = VisionEmbedder() + self.image_encoder = self.get_image_feature + self.video_encoder = self.get_video_feature + + def separate_deepstack_embeds(self, embedding: torch.Tensor): + divisor = 1 + self.num_deepstack_embeddings + if embedding.shape[-1] % divisor != 0: + raise ValueError( + f"hidden_state of {embedding.shape} should be divisible by {divisor}" + ) + separate_index = self.config.hidden_size + input_embeds = embedding[:, :separate_index] + input_deepstack_embeds = embedding[:, separate_index:] + return input_embeds, input_deepstack_embeds + + def pad_input_ids(self, input_ids: list[int], mm_inputs: MultimodalInputs): + return pad_input_tokens(input_ids, mm_inputs) + + def get_image_feature(self, items: list[MultimodalDataItem]) -> torch.Tensor: + """Eager image encode via the ``pre_encode`` / ``forward_blocks`` / + ``post_encode`` decomposition the cudagraph wrapper uses, so eager + and captured paths share a single source of truth.""" + tokens, grid = self.pre_encode(items) + metadata = self.visual.prepare_metadata(grid) + encoded = self.visual.forward_blocks(tokens, metadata) + return self.post_encode([encoded], grid) + + def get_video_feature(self, items: list[MultimodalDataItem]) -> torch.Tensor: + """Eager video encode; the cudagraph path uses the same pre/post hooks.""" + tokens, grid = self.pre_encode(items) + metadata = self.visual.prepare_metadata(grid) + encoded = self.visual.forward_blocks(tokens, metadata) + return self.post_encode([encoded], grid) + + def pre_encode( + self, + items: list[MultimodalDataItem], + ) -> tuple[torch.Tensor, torch.Tensor]: + """Eager patch-embed before the captured region; returns ``(tokens, grid)``. + + The grid field is selected per item by modality (``video_grid_thw`` for + video, ``image_grid_thw`` otherwise) so a single shared encoder cudagraph + wrapper can serve both image and video batches. + """ + device = self.visual.device + pixel_values = torch.cat( + [item.feature.to(device, non_blocking=True) for item in items], dim=0 + ).type(self.visual.dtype) + grid = torch.concat( + [ + getattr( + item, + ( + "video_grid_thw" + if item.modality == Modality.VIDEO + else "image_grid_thw" + ), + ) + for item in items + ], + dim=0, + ) + if pixel_values.dim() != 2: + raise ValueError(f"pixel_values must be 2D, got {pixel_values.dim()}D.") + if grid.dim() != 2: + raise ValueError(f"grid must be 2D, got {grid.dim()}D.") + x = self.visual.prepare_patch_embed(pixel_values, grid) + return x, grid + + def post_encode( + self, encoder_outs: list[torch.Tensor], grid: torch.Tensor + ) -> torch.Tensor: + """Eager step after the captured region; returns features.""" + return torch.cat(encoder_outs, dim=0) + + def _build_encoder_cudagraph_wrapper( + self, + mapping, + *, + max_metadata_sequences_per_batch: int | None = None, + metadata_sequence_budget_from_encoder_output_budget: bool = False, + ): + # Captured region is ``Qwen3VLMoeVisionModel.forward_blocks`` (blocks + + # deepstack mergers + merger); the merger applies a + # ``spatial_merge_size ** 2`` token reduction, so budgets count + # post-merge tokens while the capture input buffer holds + # ``spatial_merge_size ** 2 * budget`` patches. + adapter = VisionEncoderCudaGraphAdapter( + tower=self.visual, + pre_encode=self.pre_encode, + post_encode=self.post_encode, + out_div=self.visual.spatial_merge_size**2, + merge=self.visual.spatial_merge_size, + input_feature_shape=(1, self.visual.hidden_size), + modality_name="vision", + capture_tp_size=mapping.vision.tp_size, + capture_tp_group=mapping.vision.tp_group, + ) + return EncoderCudaGraphWrapper( + adapter=adapter, + budget_range=(64, 4096), + max_metadata_sequences_per_batch=max_metadata_sequences_per_batch, + metadata_sequence_budget_from_encoder_output_budget=( + metadata_sequence_budget_from_encoder_output_budget + ), + ) + + def make_encoder_cudagraph_wrappers(self, mapping): + max_video_metadata_sequences = ( + envs.TOKENSPEED_MM_VIDEO_ENCODER_CUDA_GRAPH_MAX_SEQUENCES_PER_BATCH.get() + ) + if max_video_metadata_sequences is not None: + max_video_metadata_sequences = max(1, max_video_metadata_sequences) + # Image and video encode through the identical captured region + # (``visual.forward_blocks`` over the same post-merge token buckets), so + # one wrapper serves both -- ``pre_encode`` selects the grid field per + # item by modality. Sharing a single set of budget graphs (rather than + # one set per modality) halves the captured-graph GPU memory. The video + # metadata-sequence policy is the superset (a video batch packs more + # sequences per item than an image batch at a given token budget), so it + # also covers image batches. + shared = self._build_encoder_cudagraph_wrapper( + mapping, + max_metadata_sequences_per_batch=max_video_metadata_sequences, + metadata_sequence_budget_from_encoder_output_budget=( + max_video_metadata_sequences is None + ), + ) + return {"image_encoder": shared, "video_encoder": shared} + + def get_input_embeddings(self): + return self.model.embed_tokens + + @torch.no_grad() + def forward( + self, + ctx: ForwardContext, + input_ids: torch.Tensor, + positions: torch.Tensor, + out_cache_loc: torch.Tensor, + **kwargs, + ) -> torch.Tensor: + multimodal_context = kwargs.pop("multimodal_context", None) + if ( + multimodal_context is None + or not multimodal_context.has_extend_inputs() + or ctx.forward_mode.is_decode_or_idle() + ): + return super().forward( + ctx, + input_ids, + positions, + out_cache_loc, + **kwargs, + ) + + input_embeds, model_kwargs = self.vision_embedder.apply( + input_ids=input_ids, + text_embedding=self.model.get_input_embeddings(), + ctx=multimodal_context, + encoders={ + Modality.IMAGE: EncoderSpec(self.image_encoder, deepstack=True), + Modality.VIDEO: EncoderSpec(self.video_encoder, deepstack=True), + }, + multimodal_model=self, + is_decode_or_idle=ctx.forward_mode.is_decode_or_idle(), + ) + hidden_states, aux_hidden_states = self.model( + input_ids, + positions, + ctx, + out_cache_loc, + input_embeds=input_embeds, + **model_kwargs, + ) + logits_metadata = LogitsMetadata.from_forward_context(ctx) + return self.logits_processor( + input_ids, + hidden_states, + self.lm_head, + logits_metadata, + aux_hidden_states, + ) + + def resolve_model( + self, + config: Qwen3_5TextConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None, + prefix: str, + ): + return self.model_cls( + config=config, + mapping=mapping, + quant_config=quant_config, + prefix=add_prefix("model.language_model", prefix), + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + # GDN (GatedDeltaNet) linear attention projections + # Split checkpoint format (separate qkv/z/b/a files) + ("in_proj_qkvzba.", "in_proj_qkv.", (0, 1, 2)), + ("in_proj_qkvzba.", "in_proj_z.", 3), + ("in_proj_qkvzba.", "in_proj_b.", 4), + ("in_proj_qkvzba.", "in_proj_a.", 5), + # Pre-packed checkpoint format (already merged qkvz and ba) + ("in_proj_qkvzba.", "in_proj_qkvz.", (0, 1, 2, 3)), + ("in_proj_qkvzba.", "in_proj_ba.", (4, 5)), + ] + + loaded_params: set[str] = set() + params_dict = dict(self.named_parameters(remove_duplicate=False)) + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + if "mtp" in name: + continue + if not self.is_multimodal_active and "visual" in name: + continue + # Vision-only role: drop every non-visual (LM / lm_head / norm / + # embed) weight up front, before any rename or params_dict lookup, + # so none is routed into a None module. self.model is None here, so + # named_parameters() exposes only visual params. + if getattr(self, "encoder_only", False) and "visual" not in name: + continue + if "language_model" in name: + name = name.replace(r"model.language_model.", r"model.") + if ".self_attn." in name: + name = name.replace(".self_attn", "") + if "visual" in name: + name = name.replace(r"attn.qkv.", r"attn.qkv_proj.") + name = name.replace(r"model.visual.", r"visual.") + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if "visual" in name or "mlp.experts" in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader") + weight_loader(param, loaded_weight, shard_id) + break + else: + if "visual" in name: + name = name.replace(r"attn.qkv.", r"attn.qkv_proj.") + name = name.replace(r"model.visual.", r"visual.") + # Skip loading extra bias for GPTQ models. + if name.endswith(".bias") and name not in params_dict: + continue + if name not in params_dict: + logger.warning("Parameter %s not found in params_dict", name) + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params + + +class Qwen3_5MoeForConditionalGeneration(Qwen3_5ForConditionalGeneration): + """Qwen3.5 MoE Vision-Language Model.""" + + model_cls = Qwen3_5MoeForCausalLM + + def __init__( + self, + config: Qwen3_5Config, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + is_multimodal_active: bool = True, + mm_attention_backend: str | None = None, + ) -> None: + super().__init__( + config=config, + mapping=mapping, + quant_config=quant_config, + prefix=prefix, + is_multimodal_active=is_multimodal_active, + mm_attention_backend=mm_attention_backend, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + # GDN (GatedDeltaNet) linear attention projections + # Split checkpoint format (separate qkv/z/b/a files) + ("in_proj_qkvzba.", "in_proj_qkv.", (0, 1, 2)), + ("in_proj_qkvzba.", "in_proj_z.", 3), + ("in_proj_qkvzba.", "in_proj_b.", 4), + ("in_proj_qkvzba.", "in_proj_a.", 5), + # Pre-packed checkpoint format (already merged qkvz and ba) + ("in_proj_qkvzba.", "in_proj_qkvz.", (0, 1, 2, 3)), + ("in_proj_qkvzba.", "in_proj_ba.", (4, 5)), + ] + + ignore_suffixes = ( + ".bias", + "_bias", + ".k_scale", + "_k_scale", + ".v_scale", + "_v_scale", + ".weight_scale", + "_weight_scale", + ".input_scale", + "_input_scale", + ) + loaded_params: set[str] = set() + params_dict = dict(self.named_parameters(remove_duplicate=False)) + # MoE expert weights, scales, and activation scales are handled + # by the checkpoint loader. + moe_loader = build_moe_checkpoint_loader( + params_dict=params_dict, + expert_schema=ExpertCheckpointSchema( + gate_proj_name="gate_proj", + down_proj_name="down_proj", + up_proj_name="up_proj", + ), + fused_schema=ExpertCheckpointSchema( + gate_up_fused_name="gate_up_proj", + down_proj_name="down_proj", + ), + num_experts=self.config.num_experts, + ep_rank=self.mapping.moe.ep_rank, + ep_size=self.mapping.moe.ep_size, + ) + + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + if "mtp" in name: + continue + if not self.is_multimodal_active and "visual" in name: + continue + # Vision-only role: drop every non-visual (LM / lm_head / norm / + # embed / expert) weight up front, before any rename, params_dict + # lookup, or moe_loader.load (which would KeyError on a missing + # expert param). self.model is None here, so named_parameters() + # exposes only visual params. + if getattr(self, "encoder_only", False) and "visual" not in name: + continue + if "language_model" in name: + name = name.replace(r"model.language_model.", r"model.") + if ".self_attn." in name: + name = name.replace(".self_attn", "") + if "visual" in name: + name = name.replace(r"attn.qkv.", r"attn.qkv_proj.") + name = name.replace(r"model.visual.", r"visual.") + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if "visual" in name: + continue + if "mlp.experts" in name: + continue + name = name.replace(weight_name, param_name) + # Skip loading extra parameters for GPTQ/nvfp4 models. + if name.endswith(ignore_suffixes) and name not in params_dict: + continue + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith((".bias", "_bias")) and name not in params_dict: + continue + if moe_loader.matches(name): + mapped_name = moe_loader.load(name, loaded_weight) + loaded_params.add(mapped_name) + continue + if moe_loader.is_expert_checkpoint_weight(name): + continue + + # Skip loading extra parameters for GPTQ/nvfp4 models. + if name.endswith(ignore_suffixes) and name not in params_dict: + continue + + if name in params_dict.keys(): + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + else: + logger.warning("Parameter %s not found in params_dict", name) + loaded_params.add(name) + + return loaded_params + + @classmethod + def get_model_config_for_expert_location(cls, config): + text_config = getattr(config, "text_config", config) + return ModelConfigForExpertLocation( + num_layers=text_config.num_hidden_layers, + num_logical_experts=text_config.num_experts, + num_groups=None, + ) + + +@triton.jit +def fused_qkvzba_split_reshape_cat_contiguous_kernel( + mixed_qkv, + z, + b, + a, + mixed_qkvz, + mixed_ba, + stride_qkvz, + stride_ba, + NUM_HEADS_QK: tl.constexpr, + NUM_HEADS_V: tl.constexpr, + HEAD_QK: tl.constexpr, + HEAD_V: tl.constexpr, +): + i_bs, i_qk = tl.program_id(0), tl.program_id(1) + + V_PER_GROUP: tl.constexpr = NUM_HEADS_V // NUM_HEADS_QK + + # ── Input dimensions ── + TOTAL_Q: tl.constexpr = NUM_HEADS_QK * HEAD_QK + TOTAL_K: tl.constexpr = NUM_HEADS_QK * HEAD_QK + TOTAL_V: tl.constexpr = NUM_HEADS_V * HEAD_V + + # ── Output dimensions ── + QKV_DIM_T: tl.constexpr = TOTAL_Q + TOTAL_K + TOTAL_V + + # ── Read from input (supports non-contiguous stride) ── + # q for head group i_qk: in the all_q region, offset i_qk * HEAD_QK + blk_q_ptr = mixed_qkvz + i_bs * stride_qkvz + i_qk * HEAD_QK + tl.arange(0, HEAD_QK) + # k for head group i_qk: in the all_k region + blk_k_ptr = ( + mixed_qkvz + + i_bs * stride_qkvz + + TOTAL_Q + + i_qk * HEAD_QK + + tl.arange(0, HEAD_QK) + ) + # v for head group i_qk: in the all_v region + blk_v_ptr = ( + mixed_qkvz + + i_bs * stride_qkvz + + TOTAL_Q + + TOTAL_K + + i_qk * V_PER_GROUP * HEAD_V + + tl.arange(0, V_PER_GROUP * HEAD_V) + ) + # z for head group i_qk: in the all_z region + blk_z_ptr = ( + mixed_qkvz + + i_bs * stride_qkvz + + TOTAL_Q + + TOTAL_K + + TOTAL_V + + i_qk * V_PER_GROUP * HEAD_V + + tl.arange(0, V_PER_GROUP * HEAD_V) + ) + + # ── Write to output (identical layout to the interleaved kernel) ── + blk_q_st_ptr = mixed_qkv + i_bs * QKV_DIM_T + i_qk * HEAD_QK + tl.arange(0, HEAD_QK) + blk_k_st_ptr = ( + mixed_qkv + + i_bs * QKV_DIM_T + + NUM_HEADS_QK * HEAD_QK + + i_qk * HEAD_QK + + tl.arange(0, HEAD_QK) + ) + blk_v_st_ptr = ( + mixed_qkv + + i_bs * QKV_DIM_T + + NUM_HEADS_QK * HEAD_QK * 2 + + i_qk * V_PER_GROUP * HEAD_V + + tl.arange(0, V_PER_GROUP * HEAD_V) + ) + blk_z_st_ptr = ( + z + + i_bs * NUM_HEADS_V * HEAD_V + + i_qk * V_PER_GROUP * HEAD_V + + tl.arange(0, V_PER_GROUP * HEAD_V) + ) + + tl.store(blk_q_st_ptr, tl.load(blk_q_ptr)) + tl.store(blk_k_st_ptr, tl.load(blk_k_ptr)) + tl.store(blk_v_st_ptr, tl.load(blk_v_ptr)) + tl.store(blk_z_st_ptr, tl.load(blk_z_ptr)) + + # ── b and a ── + for i in tl.static_range(V_PER_GROUP): + blk_b_ptr = mixed_ba + i_bs * stride_ba + i_qk * V_PER_GROUP + i + blk_b_st_ptr = b + i_bs * NUM_HEADS_V + i_qk * V_PER_GROUP + i + tl.store(blk_b_st_ptr, tl.load(blk_b_ptr)) + + for i in tl.static_range(V_PER_GROUP): + blk_a_ptr = mixed_ba + i_bs * stride_ba + NUM_HEADS_V + i_qk * V_PER_GROUP + i + blk_a_st_ptr = a + i_bs * NUM_HEADS_V + i_qk * V_PER_GROUP + i + tl.store(blk_a_st_ptr, tl.load(blk_a_ptr)) + + +def fused_qkvzba_split_reshape_cat_contiguous( + mixed_qkvz, + mixed_ba, + num_heads_qk, + num_heads_v, + head_qk, + head_v, +): + """Fused split/reshape/cat for Qwen3.5. Supports non-contiguous inputs. + + Input layout (per row): + mixed_qkvz: [all_q | all_k | all_v | all_z] + mixed_ba: [all_b | all_a] + + Output layout: + mixed_qkv: [all_q | all_k | all_v] (z stripped) + z: [num_v_heads, head_v] + b: [num_v_heads] + a: [num_v_heads] + """ + batch, seq_len = mixed_qkvz.shape[0], 1 + qkv_dim_t = num_heads_qk * head_qk * 2 + num_heads_v * head_v + mixed_qkv = torch.empty( + [batch * seq_len, qkv_dim_t], + dtype=mixed_qkvz.dtype, + device=mixed_qkvz.device, + ) + z = torch.empty( + [batch * seq_len, num_heads_v, head_v], + dtype=mixed_qkvz.dtype, + device=mixed_qkvz.device, + ) + b = torch.empty( + [batch * seq_len, num_heads_v], + dtype=mixed_ba.dtype, + device=mixed_ba.device, + ) + a = torch.empty_like(b) + grid = (batch * seq_len, num_heads_qk) + fused_qkvzba_split_reshape_cat_contiguous_kernel[grid]( + mixed_qkv, + z, + b, + a, + mixed_qkvz, + mixed_ba, + mixed_qkvz.stride(0), + mixed_ba.stride(0), + num_heads_qk, + num_heads_v, + head_qk, + head_v, + num_warps=1, + num_stages=3, + ) + return mixed_qkv, z, b, a + + +EntryClass = [Qwen3_5MoeForConditionalGeneration, Qwen3_5ForConditionalGeneration] diff --git a/python/tokenspeed/runtime/models/qwen3_5_moe.py b/python/tokenspeed/runtime/models/qwen3_5_moe.py new file mode 100755 index 0000000..44616d4 --- /dev/null +++ b/python/tokenspeed/runtime/models/qwen3_5_moe.py @@ -0,0 +1,392 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + + +"""Qwen3.5 MoE blocks shared by dense and MoE model variants.""" + +from __future__ import annotations + +import torch +from tokenspeed_kernel.ops.activation.triton import fused_gate_sigmoid_mul_add +from tokenspeed_kernel.ops.gemm.cute_dsl import ( + nvfp4_gemm_swiglu_nvfp4_quant, +) +from tokenspeed_kernel.ops.quantization.flashinfer import fp4_quantize +from tokenspeed_kernel.platform import current_platform +from torch import nn + +from tokenspeed.runtime.configs.qwen3_5_text_base_config import Qwen3_5BaseTextConfig +from tokenspeed.runtime.distributed.comm_manager import CommManager +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.execution.cuda_graph_wrapper import get_is_capture_mode +from tokenspeed.runtime.layers.activation import SiluAndMul +from tokenspeed.runtime.layers.dense.nvfp4 import Nvfp4LinearMethod +from tokenspeed.runtime.layers.linear import ( + MergedColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from tokenspeed.runtime.layers.moe.expert import MoELayer +from tokenspeed.runtime.layers.moe.topk import TopK +from tokenspeed.runtime.layers.moe.utils import ( + RoutingMethodType, + get_all2all_backend, + get_moe_backend, +) +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.utils import add_prefix +from tokenspeed.runtime.utils.cuda_stream import StreamFork +from tokenspeed.runtime.utils.env import envs, global_server_args_dict +from tokenspeed.runtime.utils.pdl import pdl_enabled + +_is_blackwell = current_platform().is_blackwell + + +def _is_moe_layer(layer_id: int, config) -> bool: + """Return whether the given decoder layer should use the MoE block.""" + if layer_id < 0: + return False + mlp_only_layers = getattr(config, "mlp_only_layers", []) + if layer_id in mlp_only_layers: + return False + return config.num_experts > 0 and (layer_id + 1) % config.decoder_sparse_step == 0 + + +class Qwen3_5MoeMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, + prefix: str = "", + ) -> None: + super().__init__() + self.mapping = mapping + if mapping.dense.has_tp: + tp_size = mapping.dense.tp_size + tp_rank = mapping.dense.tp_rank + tp_group = mapping.dense.tp_group + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + tp_size=tp_size, + tp_rank=tp_rank, + tp_group=tp_group, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + tp_size=tp_size, + tp_rank=tp_rank, + tp_group=tp_group, + reduce_results=reduce_results, + quant_config=quant_config, + prefix=add_prefix("down_proj", prefix), + ) + else: + self.gate_up_proj = ReplicatedLinear( + hidden_size, + intermediate_size * 2, + bias=False, + quant_config=quant_config, + prefix=add_prefix("gate_up_proj", prefix), + ) + self.down_proj = ReplicatedLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + prefix=add_prefix("down_proj", prefix), + ) + + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. " + "Only silu is supported for now." + ) + self.act_fn = SiluAndMul() + + self._use_nvfp4_gemm_swiglu_nvfp4_quant = ( + envs.TOKENSPEED_NVFP4_GEMM_SWIGLU_NVFP4_QUANT.get() + and _is_blackwell + and isinstance(self.gate_up_proj.quant_method, Nvfp4LinearMethod) + and isinstance(self.down_proj.quant_method, Nvfp4LinearMethod) + ) + self.gate_up_proj.interleave_linear_and_gate = ( + self._use_nvfp4_gemm_swiglu_nvfp4_quant + ) + + def forward(self, x): + if x.shape[0] == 0: + return x + if self._use_nvfp4_gemm_swiglu_nvfp4_quant: + x_fc1_fp4, x_fc1_scale = fp4_quantize( + x, + self.gate_up_proj.input_scale_inv, + enable_pdl=pdl_enabled(), + ) + x_fp4, x_scale = nvfp4_gemm_swiglu_nvfp4_quant( + x_fc1_fp4, + x_fc1_scale, + self.gate_up_proj.weight_swiglu_interleaved, + self.gate_up_proj.weight_scale_swiglu_interleaved, + self.gate_up_proj.alpha, + self.down_proj.input_scale_inv, + enable_pdl=pdl_enabled(), + ) + x, _ = self.down_proj((x_fp4, x_scale)) + return x + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class Qwen3_5MoeSparseMoeBlock(nn.Module): + def __init__( + self, + config: Qwen3_5BaseTextConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + layer_index: int = -1, + prefix: str = "", + alt_stream: torch.cuda.Stream | None = None, + ): + super().__init__() + self.mapping = mapping + self.layer_index = layer_index + self.tp_size = mapping.world_size + self.stream_fork = StreamFork(alt_stream) + # DeepEP is only supported with the nvfp4 cutedsl MoE backend. + # Draft models (non-quantized) must fall back to the TP path even + # when the target model has deep_ep configured globally. + self.use_deepep = ( + get_all2all_backend().is_deepep() + and get_moe_backend().is_flashinfer_cutedsl() + ) + self.comm_manager = CommManager( + mapping=mapping, + layer_id=layer_index, + is_moe=True, + prev_is_moe=_is_moe_layer(layer_index - 1, config), + ) + + if self.tp_size > config.num_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than " + f"the number of experts {config.num_experts}." + ) + + self.gate = ReplicatedLinear( + config.hidden_size, + config.num_experts, + bias=False, + quant_config=None, + prefix=add_prefix("gate", prefix), + ) + self.experts = MoELayer( + top_k=config.num_experts_per_tok, + num_experts=config.num_experts + + global_server_args_dict["ep_num_redundant_experts"], + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + quant_config=quant_config, + layer_index=layer_index, + prefix=prefix, + tp_rank=self.mapping.moe.tp_rank, + tp_size=self.mapping.moe.tp_size, + ep_rank=self.mapping.moe.ep_rank, + ep_size=self.mapping.moe.ep_size, + routing_config={ + "routing_method_type": RoutingMethodType.RenormalizeNaive, + "normalize_topk_weights": config.norm_topk_prob, + }, + ) + self.topk = TopK( + top_k=config.num_experts_per_tok, + renormalize=config.norm_topk_prob, + use_grouped_topk=False, + output_format=self.experts.topk_output_format, + ) + if getattr(config, "shared_expert_intermediate_size", 0) > 0: + self.shared_expert = Qwen3_5MoeMLP( + hidden_size=config.hidden_size, + intermediate_size=config.shared_expert_intermediate_size, + hidden_act=config.hidden_act, + mapping=self.mapping, + quant_config=quant_config, + reduce_results=False, + prefix=add_prefix("shared_expert", prefix), + ) + self.shared_expert_gate = torch.nn.Linear(config.hidden_size, 1, bias=False) + else: + self.shared_expert = None + self.shared_expert_gate = None + + def get_moe_routed_weights(self): + """Return routed expert weights excluding auxiliary shared parameters.""" + return [ + x.data + for name, x in self.experts.named_parameters() + if name not in ["correction_bias"] and "shared_experts" not in name + ] + + def forward( + self, + hidden_states: torch.Tensor, + num_global_tokens: int, + max_num_tokens_per_gpu: int, + ctx: ForwardContext, + ) -> torch.Tensor: + if self.use_deepep: + return self._forward_deepep( + hidden_states, num_global_tokens, max_num_tokens_per_gpu, ctx + ) + return self._forward_tp( + hidden_states, num_global_tokens, max_num_tokens_per_gpu, ctx + ) + + def _forward_tp( + self, + hidden_states: torch.Tensor, + num_global_tokens: int, + max_num_tokens_per_gpu: int, + ctx: ForwardContext, + ) -> torch.Tensor: + num_tokens, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + + # Gate on local (pre-comm) tokens + router_logits, _ = self.gate(hidden_states) + + # All-gather hidden_states and router_logits for topk + experts + hidden_states = self.comm_manager.pre_mlp_comm(hidden_states, ctx) + router_logits = self.comm_manager.pre_mlp_comm(router_logits, ctx) + + shared_output = None + with self.stream_fork.scope( + enable=( + self.shared_expert is not None + and hidden_states.shape[0] > 0 + and get_is_capture_mode() + ) + ) as fork: + with fork.branch(): + if self.shared_expert is not None: + shared_output = self.shared_expert(hidden_states) + + if hidden_states.shape[0] > 0: + topk_output = self.topk(hidden_states, router_logits) + else: + topk_output = self.topk.empty_topk_output( + hidden_states.device, + hidden_states=hidden_states, + router_logits=router_logits, + ) + + final_hidden_states = self.experts( + hidden_states=hidden_states, + topk_output=topk_output, + num_global_tokens=num_global_tokens, + max_num_tokens_per_gpu=max_num_tokens_per_gpu, + ) + + if shared_output is not None: + if self.shared_expert_gate is not None and hidden_states.shape[0] > 0: + fused_gate_sigmoid_mul_add( + hidden_states, + self.shared_expert_gate.weight.squeeze(0), + shared_output, + final_hidden_states, + ) + else: + final_hidden_states = final_hidden_states + shared_output + + # Reduce-scatter / all-reduce expert output back to local token count + final_hidden_states, _ = self.comm_manager.post_mlp_fused( + final_hidden_states, None, ctx + ) + + return final_hidden_states.view(num_tokens, hidden_dim) + + def _forward_deepep( + self, + hidden_states: torch.Tensor, + num_global_tokens: int, + max_num_tokens_per_gpu: int, + ctx: ForwardContext, + ) -> torch.Tensor: + """DeepEP path: routing on local tokens, dispatch/combine handled by executor.""" + num_tokens, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + + # Gate on local tokens (no all-gather needed) + router_logits, _ = self.gate(hidden_states) + + # Shared expert on local tokens (TP-parallel, needs explicit reduce) + shared_output = None + if self.shared_expert is not None: + shared_output = self.shared_expert(hidden_states) + if self.mapping.dense.has_tp: + from tokenspeed.runtime.distributed.comm_ops import all_reduce + + shared_output = all_reduce( + shared_output, + self.mapping.dense.tp_group, + ) + + # TopK on local tokens + if hidden_states.shape[0] > 0: + topk_output = self.topk(hidden_states, router_logits) + else: + topk_output = self.topk.empty_topk_output( + hidden_states.device, + hidden_states=hidden_states, + router_logits=router_logits, + ) + + # DeepEP executor handles dispatch -> MoE GEMM -> combine internally + final_hidden_states = self.experts( + hidden_states=hidden_states, + topk_output=topk_output, + num_global_tokens=num_global_tokens, + max_num_tokens_per_gpu=max_num_tokens_per_gpu, + ) + + if shared_output is not None: + if self.shared_expert_gate is not None and hidden_states.shape[0] > 0: + fused_gate_sigmoid_mul_add( + hidden_states, + self.shared_expert_gate.weight.squeeze(0), + shared_output, + final_hidden_states, + ) + else: + final_hidden_states = final_hidden_states + shared_output + + return final_hidden_states.view(num_tokens, hidden_dim) diff --git a/python/tokenspeed/runtime/models/qwen3_5_nextn.py b/python/tokenspeed/runtime/models/qwen3_5_nextn.py new file mode 100644 index 0000000..7293136 --- /dev/null +++ b/python/tokenspeed/runtime/models/qwen3_5_nextn.py @@ -0,0 +1,419 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import logging +from collections.abc import Iterable +from dataclasses import replace + +import torch +from tokenspeed_kernel.ops.activation.triton import sigmoid_mul +from torch import nn +from transformers import PretrainedConfig + +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.execution.context import ( + ForwardContext, + report_collective_sizing, +) +from tokenspeed.runtime.execution.forward_batch_info import ForwardMode +from tokenspeed.runtime.layers.layernorm import GemmaRMSNorm +from tokenspeed.runtime.layers.linear import ReplicatedLinear +from tokenspeed.runtime.layers.logits_processor import LogitsMetadata, LogitsProcessor +from tokenspeed.runtime.layers.moe import ( + ExpertCheckpointSchema, + build_moe_checkpoint_loader, +) +from tokenspeed.runtime.layers.vocab_parallel_embedding import ParallelLMHead +from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader +from tokenspeed.runtime.models.qwen3_5 import ( + Qwen3_5AttentionDecoderLayer, + Qwen3_5ForCausalLM, +) +from tokenspeed.runtime.utils import add_prefix + +logger = logging.getLogger(__name__) + + +class Qwen3_5DraftAttentionDecoderLayer(Qwen3_5AttentionDecoderLayer): + """NextN draft variant: skip dead catch-up rows on the first draft step. + + On the first draft step the backend runs in DECODE mode with ``q`` sliced + to ``bs`` while ``self.attn`` still writes the full ``N`` rope-d KV rows + from the just-drafted tokens. Multi-step decode delegates to base. + + MIXED catch-up requires a backend that populates a decode-slot metadata + under EXTEND/MIXED at draft init (e.g. trtllm-mha); MHA-family backends + that assert ``not is_mixed()`` at metadata init are not supported. + """ + + def _attn( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gate: torch.Tensor | None, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + ) -> torch.Tensor: + if ctx.accept_lengths is None: + return super()._attn(q, k, v, gate, ctx, out_cache_loc) + + self._apply_correction(ctx) + q = q.index_select(0, ctx.gather_ids) + if gate is not None: + gate = gate.index_select(0, ctx.gather_ids) + # Dispatch as DECODE over the sliced live rows via self.attn (see the + # class docstring), which keeps the standard k/v reshape and KV write. + # A ctx copy overrides only the forward mode; record_kv_cache (keyed off + # the real mode) forces the backend's PD layerwise cache-step record that + # DECODE would otherwise skip on an EXTEND/MIXED catch-up. + decode_ctx = replace(ctx, forward_mode=ForwardMode.DECODE) + attn_output = self.attn( + q, + k, + v, + decode_ctx, + out_cache_loc, + record_kv_cache=not ctx.forward_mode.is_decode_or_idle(), + ) + if gate is not None: + sigmoid_mul(attn_output, gate) + return attn_output + + def _apply_correction(self, ctx: ForwardContext) -> None: + """Trim decode rows' cache_seqlens by ``spec_num_tokens - accept_lengths``.""" + seq_lens_buf = ctx.draft_seq_lens_buf + if seq_lens_buf is None or ctx.accept_lengths is None: + return + num_extends = ctx.num_extends + if num_extends >= ctx.bs: + return + correction = ( + ctx.attn_backend.spec_num_tokens - ctx.accept_lengths[num_extends:] + ).to(seq_lens_buf.dtype) + seq_lens_buf[num_extends : ctx.bs].sub_(correction).clamp_(min=1) + + def _maybe_narrow_residual( + self, + residual: torch.Tensor, + ctx: ForwardContext, + ) -> torch.Tensor: + if ctx.accept_lengths is None or ctx.forward_mode.is_idle(): + return residual + return residual.index_select(0, ctx.gather_ids) + + +class Qwen3_5DraftForCausalLM(Qwen3_5ForCausalLM): + """Causal LM with the draft-variant attention layer injected. + + Restricted to single-layer drafts: ``_apply_correction`` mutates + ``ctx.draft_seq_lens_buf`` in place and is not idempotent across layers. + A multi-layer draft would double-trim cache_seqlens. Lift the correction + out of the per-layer hook (e.g. into the drafter) before relaxing this. + """ + + ATTENTION_LAYER_CLS: type = Qwen3_5DraftAttentionDecoderLayer + + def __init__( + self, + config, + mapping, + quant_config=None, + prefix: str = "", + ) -> None: + if config.num_hidden_layers != 1: + raise ValueError( + "Qwen3_5DraftForCausalLM requires num_hidden_layers == 1 " + f"(got {config.num_hidden_layers}); _apply_correction is not " + "idempotent across layers." + ) + super().__init__(config, mapping, quant_config=quant_config, prefix=prefix) + + +class Qwen3_5ForConditionalGenerationNextN(nn.Module): + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + quant_config=None, + prefix: str = "", + ) -> None: + nn.Module.__init__(self) + + self.is_multimodal = hasattr(config, "text_config") + if self.is_multimodal: + config = config.text_config + + # The MTP model is unquantized in the nvfp4 checkpoint. + if quant_config and quant_config.get_name() == "nvfp4": + quant_config = None + + self.config = config + self.mapping = mapping + self.quant_config = quant_config + + self.fc = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False) + RMSNorm_cls = GemmaRMSNorm + self.pre_fc_norm_embedding = RMSNorm_cls( + config.hidden_size, config.rms_norm_eps + ) + self.pre_fc_norm_hidden = RMSNorm_cls(config.hidden_size, config.rms_norm_eps) + config.num_hidden_layers = 1 + config.full_attention_interval = 1 + self.model = Qwen3_5DraftForCausalLM( + config, + mapping=self.mapping, + quant_config=quant_config, + prefix=add_prefix("mtp", prefix), + ) + + if config.tie_word_embeddings: + self.lm_head = self.model.embed_tokens + else: + if self.mapping.attn.has_dp: + self.lm_head = ReplicatedLinear( + config.hidden_size, + config.vocab_size, + bias=False, + prefix=add_prefix("lm_head", prefix), + ) + else: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + prefix=add_prefix("lm_head", prefix), + ) + + self.logits_processor = LogitsProcessor( + config, + skip_all_gather=self.mapping.attn.has_dp, + tp_rank=self.mapping.attn.tp_rank, + tp_size=self.mapping.attn.tp_size, + tp_group=self.mapping.attn.tp_group, + ) + + def get_hot_token_id(self): + return None + + def get_embed_and_head(self): + return self.model.embed_tokens.weight, self.lm_head.weight + + def set_embed_and_head(self, embed, head): + del self.model.embed_tokens.weight + if not self.config.tie_word_embeddings: + del self.lm_head.weight + + self.model.embed_tokens.weight = embed + self.lm_head.weight = head + torch.cuda.empty_cache() + torch.cuda.synchronize() + + @torch.no_grad() + def forward( + self, + ctx: ForwardContext, + input_ids: torch.Tensor, + positions: torch.Tensor, + out_cache_loc: torch.Tensor, + input_embeds: torch.Tensor | None = None, + captured_hidden_states: torch.Tensor | None = None, + **kwargs, + ): + if captured_hidden_states is None and not ctx.forward_mode.is_idle(): + raise ValueError("Qwen3.5 MTP requires captured_hidden_states.") + + if ctx.forward_mode.is_idle(): + # IDLE forward: skip MTP-specific ops, just run the inner model + # for NCCL collective participation. + hidden_states = torch.zeros( + 0, + self.config.hidden_size * 2, + device=input_ids.device, + dtype=self.model.embed_tokens.weight.dtype, + ) + else: + if input_embeds is not None: + raise ValueError("input_embeds is not supported for nextn forward.") + input_embeds = self.model.embed_tokens(input_ids) + hidden_states = captured_hidden_states + input_embeds = self.pre_fc_norm_embedding(input_embeds) + hidden_states = self.pre_fc_norm_hidden(hidden_states) + hidden_states = torch.cat([input_embeds, hidden_states], dim=-1) + + hidden_states = self.fc(hidden_states) + + with report_collective_sizing(ctx, ctx.bs, ctx.global_bs): + hidden_states, _ = self.model( + input_ids, + positions, + ctx, + out_cache_loc, + input_embeds=hidden_states, + ) + + logits_metadata = LogitsMetadata.from_forward_context(ctx) + return self.logits_processor( + input_ids, hidden_states, self.lm_head, logits_metadata + ) + + def load_weights( + self, weights: Iterable[tuple[str, torch.Tensor]], is_mtp: bool = False + ): + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + num_experts = getattr(self.config, "num_experts", None) + + # Skip loading extra parameters for GPTQ/nvfp4 models. + ignore_suffixes = ( + ".bias", + "_bias", + ".k_scale", + "_k_scale", + ".v_scale", + "_v_scale", + ".weight_scale", + "_weight_scale", + ".input_scale", + "_input_scale", + ) + params_dict = dict(self.named_parameters(remove_duplicate=False)) + moe_loader = None + if num_experts is not None: + # MoE expert weights, scales, and activation scales are handled + # by the checkpoint loader. + moe_loader = build_moe_checkpoint_loader( + params_dict=params_dict, + expert_schema=ExpertCheckpointSchema( + gate_proj_name="gate_proj", + down_proj_name="down_proj", + up_proj_name="up_proj", + ), + fused_schema=ExpertCheckpointSchema( + gate_up_fused_name="gate_up_proj", + down_proj_name="down_proj", + ), + num_experts=num_experts, + ep_rank=self.mapping.moe.ep_rank, + ep_size=self.mapping.moe.ep_size, + ) + loaded_params: set[str] = set() + + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + + # Only process MTP branch weights + if "mtp" not in name: + continue + + if name.startswith("mtp."): + # Remove the mtp. prefix for processing + name = name.replace("mtp.", "model.") + + name = name.replace("model.fc", "fc") + name = name.replace("model.pre_fc", "pre_fc") + + if ".self_attn." in name: + name = name.replace(".self_attn", "") + + # 1) Process stacked parameters (q_proj/k_proj/v_proj & gate_proj/up_proj) + for param_name, weight_name, shard_id in stacked_params_mapping: + # Skip non-matching weights + if weight_name not in name: + continue + + # Skip MoE experts.* here, handled separately below + if "mlp.experts" in name: + continue + + name_mapped = name.replace(weight_name, param_name) + + # Skip loading extra parameters for GPTQ/nvfp4 models. + if ( + name_mapped.endswith(ignore_suffixes) + and name_mapped not in params_dict + ): + continue + + if name_mapped not in params_dict: + continue + + param = params_dict[name_mapped] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight, shard_id) + name = name_mapped + break + else: + # Skip loading extra bias for GPTQ models. + if name.endswith((".bias", "_bias")) and name not in params_dict: + continue + if moe_loader is not None: + if moe_loader.matches(name): + mapped_name = moe_loader.load(name, loaded_weight) + loaded_params.add(mapped_name) + continue + if moe_loader.is_expert_checkpoint_weight(name): + continue + + # Skip loading extra parameters for GPTQ/nvfp4 models. + if name.endswith(ignore_suffixes) and name not in params_dict: + continue + + if name not in params_dict: + logger.warning("MTP weight not in params_dict: %s", name) + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + loaded_params.add(name) + return loaded_params + + +class Qwen3_5MoeForConditionalGenerationNextN(Qwen3_5ForConditionalGenerationNextN): + def __init__( + self, + config: PretrainedConfig, + mapping: Mapping, + quant_config=None, + prefix: str = "", + ) -> None: + super().__init__( + config=config, mapping=mapping, quant_config=quant_config, prefix=prefix + ) + + +EntryClass = [ + Qwen3_5ForConditionalGenerationNextN, + Qwen3_5MoeForConditionalGenerationNextN, +] diff --git a/python/tokenspeed/runtime/models/qwen3_asr.py b/python/tokenspeed/runtime/models/qwen3_asr.py new file mode 100644 index 0000000..cd48720 --- /dev/null +++ b/python/tokenspeed/runtime/models/qwen3_asr.py @@ -0,0 +1,274 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Inference-only Qwen3-ASR model compatible with Hugging Face weights.""" + +from __future__ import annotations + +from collections.abc import Iterable +from typing import Any + +import torch +from torch import nn + +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader +from tokenspeed.runtime.models.qwen3 import Qwen3ForCausalLM +from tokenspeed.runtime.models.qwen3_audio import Qwen3AudioEncoder +from tokenspeed.runtime.multimodal.embedder import ( + EncoderSpec, + MultimodalEmbedder, + pad_input_tokens, +) +from tokenspeed.runtime.multimodal.inputs import ( + Modality, + MultimodalDataItem, + MultimodalInputs, +) +from tokenspeed.runtime.utils import add_prefix + + +class Qwen3ASRForConditionalGeneration(nn.Module): + """Qwen3 dense language model plus the shared Qwen audio tower.""" + + default_bitsandbytes_target_modules = [ + ".gate_proj.", + ".down_proj.", + ".up_proj.", + ".q_proj.", + ".k_proj.", + ".v_proj.", + ".o_proj.", + ] + bitsandbytes_stacked_params_mapping = { + "q_proj": ("qkv_proj", 0), + "k_proj": ("qkv_proj", 1), + "v_proj": ("qkv_proj", 2), + "gate_proj": ("gate_up_proj", 0), + "up_proj": ("gate_up_proj", 1), + } + + def __init__( + self, + config: Any, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + is_multimodal_active: bool = True, + mm_attention_backend: str | None = None, + ) -> None: + super().__init__() + self.config = config + self.mapping = mapping + self.quant_config = quant_config + self.is_multimodal_active = is_multimodal_active + + thinker_config = getattr(config, "thinker_config", config) + self.thinker_config = thinker_config + text_config = thinker_config.text_config + audio_config = getattr(thinker_config, "audio_config", None) + if audio_config is None: + raise ValueError("Qwen3-ASR config is missing thinker_config.audio_config") + if int(audio_config.output_dim) != int(text_config.hidden_size): + raise ValueError( + "audio output_dim must match the language hidden_size: " + f"{audio_config.output_dim} != {text_config.hidden_size}" + ) + + # Qwen3ForCausalLM currently owns the model/lm_head prefixes internally; + # nesting it under this attribute yields the desired runtime parameter + # names (language_model.model.*, language_model.lm_head.*). + self.language_model = Qwen3ForCausalLM( + text_config, + mapping=mapping, + quant_config=quant_config, + ) + if is_multimodal_active: + self.audio_tower = Qwen3AudioEncoder( + audio_config, + mapping, + # Keep the encoder in BF16/FP16 unless a future quantizer + # explicitly supports this tower. Text quantization remains + # active through language_model. + quant_config=None, + prefix=add_prefix("audio_tower", prefix), + mm_attention_backend=mm_attention_backend, + ) + self.multimodal_embedder = MultimodalEmbedder() + # ModelExecutor may replace this callable with an encoder graph. + self.audio_encoder = self.get_audio_feature + else: + self.audio_tower = None + self.multimodal_embedder = None + self.audio_encoder = None + + def get_audio_feature(self, items: list[MultimodalDataItem]) -> torch.Tensor: + if self.audio_tower is None: + raise RuntimeError("Qwen3-ASR audio tower is disabled") + return self.audio_tower.encode(items) + + def pad_input_ids( + self, input_ids: list[int], mm_inputs: MultimodalInputs + ) -> list[int]: + return pad_input_tokens(input_ids, mm_inputs) + + @property + def start_layer(self) -> int: + return int(getattr(self.language_model.model, "start_layer", 0)) + + @property + def end_layer(self) -> int: + return int( + getattr( + self.language_model.model, + "end_layer", + self.thinker_config.text_config.num_hidden_layers, + ) + ) + + @property + def lm_head(self): + return self.language_model.lm_head + + @property + def logits_processor(self): + return self.language_model.logits_processor + + def get_input_embeddings(self): + return self.language_model.model.embed_tokens + + def get_embed_and_head(self): + return self.language_model.get_embed_and_head() + + def set_embed_and_head(self, embed, head) -> None: + self.language_model.set_embed_and_head(embed, head) + + def load_kv_cache_scales(self, quantization_param_path: str) -> None: + self.language_model.load_kv_cache_scales(quantization_param_path) + + @torch.no_grad() + def forward( + self, + ctx, + input_ids: torch.Tensor, + positions: torch.Tensor, + out_cache_loc: torch.Tensor, + **kwargs, + ): + multimodal_context = kwargs.pop("multimodal_context", None) + if ( + multimodal_context is not None + and multimodal_context.has_extend_inputs() + and not ctx.forward_mode.is_decode_or_idle() + ): + if self.multimodal_embedder is None or self.audio_encoder is None: + raise RuntimeError( + "audio input was provided while Qwen3-ASR runs with " + "--language-model-only" + ) + input_embeds, model_kwargs = self.multimodal_embedder.apply( + input_ids=input_ids, + text_embedding=self.get_input_embeddings(), + ctx=multimodal_context, + encoders={ + Modality.AUDIO: EncoderSpec( + self.audio_encoder, + deepstack=False, + ) + }, + multimodal_model=self, + is_decode_or_idle=ctx.forward_mode.is_decode_or_idle(), + ) + kwargs.update(model_kwargs) + if input_embeds is not None: + kwargs["input_embeds"] = input_embeds + + return self.language_model( + ctx, + input_ids, + positions, + out_cache_loc, + **kwargs, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Stream the official ``thinker.*`` ASR checkpoint layout.""" + loaded: set[str] = set() + params = dict(self.named_parameters(remove_duplicate=False)) + + for name, tensor in weights: + if "talker." in name or "code2wav." in name: + continue + if name.startswith("thinker.audio_tower.") or name.startswith( + "audio_tower." + ): + if self.audio_tower is None: + continue + loaded_name = self.audio_tower.load_weight(name, tensor) + if loaded_name is None: + raise ValueError(f"unknown Qwen audio weight {name}") + loaded.add(f"audio_tower.{loaded_name}") + continue + + name, shard_id = self.map_language_weight_name(name) + if ( + self.thinker_config.text_config.tie_word_embeddings + and name == "language_model.lm_head.weight" + ): + continue + if name not in params: + raise ValueError(f"unknown Qwen3-ASR language weight {name}") + param = params[name] + if shard_id is not None: + param.weight_loader(param, tensor, shard_id) + else: + loader = getattr(param, "weight_loader", default_weight_loader) + loader(param, tensor) + loaded.add(name) + return loaded + + @staticmethod + def map_language_weight_name(name: str) -> tuple[str, str | int | None]: + """Map an official text checkpoint name to a wrapper parameter.""" + if name.startswith("thinker.model."): + name = name.replace("thinker.model.", "language_model.model.", 1) + elif name.startswith("thinker.lm_head."): + name = name.replace("thinker.lm_head.", "language_model.lm_head.", 1) + elif name.startswith("language_model."): + pass + elif name.startswith("thinker."): + raise ValueError(f"unsupported Qwen3-ASR thinker weight {name}") + else: + name = f"language_model.{name}" + + for checkpoint_name, fused_name, shard_id in ( + (".q_proj.", ".qkv_proj.", "q"), + (".k_proj.", ".qkv_proj.", "k"), + (".v_proj.", ".qkv_proj.", "v"), + (".gate_proj.", ".gate_up_proj.", 0), + (".up_proj.", ".gate_up_proj.", 1), + ): + if checkpoint_name in name: + return name.replace(checkpoint_name, fused_name), shard_id + return name, None + + +EntryClass = [Qwen3ASRForConditionalGeneration] diff --git a/python/tokenspeed/runtime/models/qwen3_audio.py b/python/tokenspeed/runtime/models/qwen3_audio.py new file mode 100644 index 0000000..2a4439d --- /dev/null +++ b/python/tokenspeed/runtime/models/qwen3_audio.py @@ -0,0 +1,559 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Shared Qwen3-ASR/Qwen3-Omni audio encoder. + +The gateway extracts log-Mel features with the Whisper-compatible Qwen3 +frontend. Each multimodal item contains one ``[num_mel_bins, num_frames]`` +tensor (a leading singleton batch dimension is also accepted) and optionally +carries ``audio_feature_lengths`` or ``feature_attention_mask`` in +``model_specific_data``. This module packs those request-local tensors and +implements the common Qwen3 audio tower used by both Qwen3-ASR and the +Qwen3-Omni thinker. +""" + +from __future__ import annotations + +from collections.abc import Iterable, Sequence +from numbers import Integral +from typing import Any + +import torch +import torch.nn.functional as F +from torch import nn +from transformers.activations import ACT2FN + +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.layers.attention.mm_encoder_attention import ( + MultimodalEncoderAttention, +) +from tokenspeed.runtime.layers.linear import ( + ColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader +from tokenspeed.runtime.multimodal.inputs import MultimodalDataItem +from tokenspeed.runtime.utils import add_prefix + + +def _cnn_output_lengths(input_lengths: torch.Tensor) -> torch.Tensor: + """Lengths after the tower's three stride-2, padding-1 convolutions.""" + lengths = input_lengths + for _ in range(3): + lengths = (lengths - 1) // 2 + 1 + return lengths + + +def qwen3_audio_output_lengths( + input_lengths: torch.Tensor | Sequence[int] | int, + *, + n_window: int = 50, +) -> torch.Tensor: + """Return language-token counts produced for log-Mel frame lengths. + + Qwen processes audio in chunks of ``2 * n_window`` frames. Convolution + padding is applied independently to every chunk, so applying a single + ``ceil(length / 8)`` to a long clip would under-count tokens at chunk + boundaries. + + Args: + input_lengths: Scalar or one-dimensional Mel-frame lengths. + n_window: Half of the frontend chunk size from the audio config. + + Returns: + A one-dimensional ``torch.long`` tensor with one token count per input. + """ + if n_window <= 0: + raise ValueError(f"n_window must be positive, got {n_window}") + if isinstance(input_lengths, Integral): + lengths = torch.tensor([int(input_lengths)], dtype=torch.long) + elif isinstance(input_lengths, torch.Tensor): + lengths = input_lengths.to(dtype=torch.long) + if lengths.ndim == 0: + lengths = lengths.unsqueeze(0) + else: + lengths = torch.as_tensor(list(input_lengths), dtype=torch.long) + if lengths.ndim != 1: + raise ValueError( + f"audio feature lengths must be one-dimensional, got {lengths.shape}" + ) + if torch.any(lengths < 0): + raise ValueError("audio feature lengths cannot be negative") + + chunk_size = 2 * n_window + full_chunks = lengths // chunk_size + tail = lengths % chunk_size + full_chunk_output = int( + _cnn_output_lengths(torch.tensor([chunk_size], dtype=torch.long)).item() + ) + tail_output = _cnn_output_lengths(tail) + tail_output = torch.where(tail == 0, torch.zeros_like(tail), tail_output) + return full_chunks * full_chunk_output + tail_output + + +def _one_item_feature_length(item: MultimodalDataItem, max_frames: int) -> int: + explicit_length = item.model_specific_data.get("audio_feature_lengths") + if explicit_length is not None: + if isinstance(explicit_length, torch.Tensor): + if explicit_length.numel() != 1: + raise ValueError( + "one audio item must carry exactly one audio_feature_lengths " + f"value, got shape {explicit_length.shape}" + ) + length = int(explicit_length.reshape(-1)[0].item()) + elif isinstance(explicit_length, Integral): + length = int(explicit_length) + else: + raise TypeError( + "audio_feature_lengths must be an integer or tensor, got " + f"{type(explicit_length).__name__}" + ) + else: + attention_mask = item.model_specific_data.get("feature_attention_mask") + if attention_mask is None: + length = max_frames + else: + attention_mask = torch.as_tensor(attention_mask) + if attention_mask.ndim == 2 and attention_mask.shape[0] == 1: + attention_mask = attention_mask.squeeze(0) + if attention_mask.ndim != 1: + raise ValueError( + "one audio item must carry a one-dimensional " + f"feature_attention_mask, got {attention_mask.shape}" + ) + if attention_mask.numel() != max_frames: + raise ValueError( + "feature_attention_mask length does not match the audio " + f"feature: {attention_mask.numel()} != {max_frames}" + ) + # Whisper masks are right-padded. Requiring a contiguous prefix + # prevents silently packing frames from a malformed sparse mask. + mask = attention_mask.to(dtype=torch.bool) + length = int(mask.sum().item()) + expected = torch.arange(max_frames, device=mask.device) < length + if not torch.equal(mask, expected): + raise ValueError("feature_attention_mask must be right-padded") + + if length <= 0 or length > max_frames: + raise ValueError( + f"audio feature length must be in [1, {max_frames}], got {length}" + ) + return length + + +def pack_qwen3_audio_features( + items: Sequence[MultimodalDataItem], + *, + num_mel_bins: int, + device: torch.device, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + """Pack itemized log-Mel tensors for :class:`Qwen3AudioEncoder`. + + Returns a feature tensor shaped ``[num_mel_bins, sum(valid_frames)]`` and + a length tensor shaped ``[num_items]``. Padding is removed before packing. + """ + if not items: + raise ValueError("at least one audio item is required") + + features: list[torch.Tensor] = [] + lengths: list[int] = [] + for item in items: + feature = item.feature + if not isinstance(feature, torch.Tensor): + raise TypeError( + "audio item feature must be materialized as a torch.Tensor, got " + f"{type(feature).__name__}" + ) + if feature.ndim == 3 and feature.shape[0] == 1: + feature = feature.squeeze(0) + if feature.ndim != 2: + raise ValueError( + "audio feature must have shape [mel, frames] or [1, mel, frames], " + f"got {feature.shape}" + ) + if feature.shape[0] != num_mel_bins: + raise ValueError( + f"expected {num_mel_bins} Mel bins, got feature shape {feature.shape}" + ) + length = _one_item_feature_length(item, feature.shape[1]) + features.append(feature[:, :length].to(device=device, dtype=dtype)) + lengths.append(length) + + return ( + torch.cat(features, dim=1).contiguous(), + torch.tensor(lengths, dtype=torch.long, device=device), + ) + + +class SinusoidsPositionEmbedding(nn.Module): + """Non-trainable absolute position embedding used by the audio tower.""" + + def __init__(self, length: int, channels: int, max_timescale: int = 10000): + super().__init__() + if channels % 2: + raise ValueError("sinusoidal position embeddings need even channels") + log_increment = torch.log(torch.tensor(float(max_timescale))) / ( + channels // 2 - 1 + ) + inv_timescales = torch.exp(-log_increment * torch.arange(channels // 2)) + positions = torch.arange(length).unsqueeze(1) * inv_timescales.unsqueeze(0) + self.register_buffer( + "positional_embedding", + torch.cat([positions.sin(), positions.cos()], dim=1), + persistent=False, + ) + + def forward(self, sequence_length: int) -> torch.Tensor: + return self.positional_embedding[:sequence_length] + + +class Qwen3AudioEncoderLayer(nn.Module): + def __init__( + self, + config: Any, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + mm_attention_backend: str | None = None, + ) -> None: + super().__init__() + self.self_attn = MultimodalEncoderAttention( + embed_dim=config.d_model, + num_heads=config.encoder_attention_heads, + mapping=mapping, + qkv_bias=True, + proj_bias=True, + quant_config=quant_config, + prefix=add_prefix("self_attn", prefix), + mm_attention_backend=mm_attention_backend, + ) + self.self_attn_layer_norm = nn.LayerNorm(config.d_model) + self.final_layer_norm = nn.LayerNorm(config.d_model) + self.activation_fn = ACT2FN[config.activation_function] + + vision = mapping.vision + self.fc1 = ColumnParallelLinear( + config.d_model, + config.encoder_ffn_dim, + bias=True, + quant_config=quant_config, + prefix=add_prefix("fc1", prefix), + tp_rank=vision.tp_rank, + tp_size=vision.tp_size, + tp_group=vision.tp_group, + ) + self.fc2 = RowParallelLinear( + config.encoder_ffn_dim, + config.d_model, + bias=True, + quant_config=quant_config, + prefix=add_prefix("fc2", prefix), + tp_rank=vision.tp_rank, + tp_size=vision.tp_size, + tp_group=vision.tp_group, + reduce_results=True, + ) + + def forward( + self, + hidden_states: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: int, + ) -> torch.Tensor: + residual = hidden_states + normed = self.self_attn_layer_norm(hidden_states) + attended = self.self_attn( + normed.unsqueeze(0), + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + ).squeeze(0) + hidden_states = residual + attended + + residual = hidden_states + hidden_states = self.final_layer_norm(hidden_states) + hidden_states, _ = self.fc1(hidden_states) + hidden_states = self.activation_fn(hidden_states) + hidden_states, _ = self.fc2(hidden_states) + hidden_states = residual + hidden_states + + if hidden_states.dtype == torch.float16: + limit = torch.finfo(hidden_states.dtype).max - 1000 + hidden_states = hidden_states.clamp(min=-limit, max=limit) + return hidden_states + + +class Qwen3AudioEncoder(nn.Module): + """Parameterizable Qwen audio tower shared by ASR and Omni thinker.""" + + @staticmethod + def _normalize_attention_backend(name: str | None) -> str | None: + # The cuDNN path consumes vision-specific sequence metadata. Audio can + # still use the platform default while an Omni vision tower uses cuDNN. + return None if name == "flashinfer_cudnn" else name + + def __init__( + self, + config: Any, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + mm_attention_backend: str | None = None, + ) -> None: + super().__init__() + self.config = config + self.num_mel_bins = int(config.num_mel_bins) + self.n_window = int(config.n_window) + self.n_window_infer = int(config.n_window_infer) + self.conv_chunksize = int(config.conv_chunksize) + mm_attention_backend = self._normalize_attention_backend(mm_attention_backend) + + self.positional_embedding = SinusoidsPositionEmbedding( + int(config.max_source_positions), int(config.d_model) + ) + self.conv2d1 = nn.Conv2d( + 1, config.downsample_hidden_size, kernel_size=3, stride=2, padding=1 + ) + self.conv2d2 = nn.Conv2d( + config.downsample_hidden_size, + config.downsample_hidden_size, + kernel_size=3, + stride=2, + padding=1, + ) + self.conv2d3 = nn.Conv2d( + config.downsample_hidden_size, + config.downsample_hidden_size, + kernel_size=3, + stride=2, + padding=1, + ) + conv_out_dim = config.downsample_hidden_size * ( + (((config.num_mel_bins + 1) // 2 + 1) // 2 + 1) // 2 + ) + self.conv_out = ReplicatedLinear( + conv_out_dim, + config.d_model, + bias=False, + quant_config=quant_config, + prefix=add_prefix("conv_out", prefix), + ) + self.layers = nn.ModuleList( + [ + Qwen3AudioEncoderLayer( + config, + mapping, + quant_config=quant_config, + prefix=add_prefix(f"layers.{index}", prefix), + mm_attention_backend=mm_attention_backend, + ) + for index in range(config.encoder_layers) + ] + ) + self.ln_post = nn.LayerNorm(config.d_model) + self.proj1 = ReplicatedLinear( + config.d_model, + config.d_model, + bias=True, + quant_config=quant_config, + prefix=add_prefix("proj1", prefix), + ) + self.act = ACT2FN[config.activation_function] + self.proj2 = ReplicatedLinear( + config.d_model, + config.output_dim, + bias=True, + quant_config=quant_config, + prefix=add_prefix("proj2", prefix), + ) + + @property + def dtype(self) -> torch.dtype: + return self.conv2d1.weight.dtype + + @property + def device(self) -> torch.device: + return self.conv2d1.weight.device + + def encode(self, items: Sequence[MultimodalDataItem]) -> torch.Tensor: + input_features, feature_lengths = pack_qwen3_audio_features( + items, + num_mel_bins=self.num_mel_bins, + device=self.device, + dtype=self.dtype, + ) + return self(input_features, feature_lengths) + + def forward( + self, + input_features: torch.Tensor, + feature_lengths: torch.Tensor, + ) -> torch.Tensor: + """Encode packed features into concatenated LM-space audio tokens. + + Args: + input_features: Packed log-Mel features shaped + ``[num_mel_bins, sum(feature_lengths)]``. + feature_lengths: Valid frame count for every original audio item. + + Returns: + Audio embeddings shaped ``[sum(output_lengths), output_dim]`` in + the same item order as ``feature_lengths``. + """ + input_features = input_features.to(device=self.device, dtype=self.dtype) + feature_lengths = feature_lengths.to(device=self.device, dtype=torch.long) + if input_features.ndim != 2 or input_features.shape[0] != self.num_mel_bins: + raise ValueError( + f"expected packed features [{self.num_mel_bins}, frames], got " + f"{input_features.shape}" + ) + if feature_lengths.ndim != 1 or feature_lengths.numel() == 0: + raise ValueError("feature_lengths must be a non-empty vector") + if torch.any(feature_lengths <= 0): + raise ValueError("feature_lengths must be positive") + if int(feature_lengths.sum().item()) != input_features.shape[1]: + raise ValueError( + "packed feature width does not match feature_lengths: " + f"{input_features.shape[1]} != {int(feature_lengths.sum().item())}" + ) + + chunk_size = self.n_window * 2 + chunk_lengths_list: list[int] = [] + for length in feature_lengths.tolist(): + full_chunks, tail = divmod(int(length), chunk_size) + chunk_lengths_list.extend([chunk_size] * full_chunks) + if tail: + chunk_lengths_list.append(tail) + chunk_lengths = torch.tensor( + chunk_lengths_list, dtype=torch.long, device=self.device + ) + + chunks = input_features.transpose(0, 1).split(chunk_lengths_list, dim=0) + padded_features = nn.utils.rnn.pad_sequence(chunks, batch_first=True).transpose( + 1, 2 + ) + chunk_output_lengths = _cnn_output_lengths(chunk_lengths) + max_chunk_output = int(chunk_output_lengths.max().item()) + output_mask = torch.arange(max_chunk_output, device=self.device).unsqueeze( + 0 + ) < chunk_output_lengths.unsqueeze(1) + + padded_features = padded_features.unsqueeze(1) + encoded_chunks: list[torch.Tensor] = [] + for conv_input in padded_features.split(self.conv_chunksize, dim=0): + hidden = F.gelu(self.conv2d1(conv_input)) + hidden = F.gelu(self.conv2d2(hidden)) + hidden = F.gelu(self.conv2d3(hidden)) + encoded_chunks.append(hidden) + hidden = torch.cat(encoded_chunks, dim=0) + + batch, channels, frequency, time = hidden.shape + hidden, _ = self.conv_out( + hidden.permute(0, 3, 1, 2) + .contiguous() + .view(batch, time, channels * frequency) + ) + if hidden.shape[1] > self.positional_embedding.positional_embedding.shape[0]: + raise ValueError( + "audio chunk exceeds max_source_positions after convolution: " + f"{hidden.shape[1]}" + ) + hidden = hidden + self.positional_embedding(hidden.shape[1]).to(hidden.dtype) + hidden_states = hidden[output_mask] + + output_lengths = qwen3_audio_output_lengths( + feature_lengths, n_window=self.n_window + ).to(device=self.device) + chunks_per_attention_window = max(1, self.n_window_infer // chunk_size) + attention_window = max_chunk_output * chunks_per_attention_window + attention_lengths: list[int] = [] + for length in output_lengths.tolist(): + full_windows, tail = divmod(int(length), attention_window) + attention_lengths.extend([attention_window] * full_windows) + if tail: + attention_lengths.append(tail) + cu_seqlens = torch.tensor( + [0, *attention_lengths], dtype=torch.int32, device=self.device + ).cumsum(0, dtype=torch.int32) + max_seqlen = max(attention_lengths) + + for layer in self.layers: + hidden_states = layer(hidden_states, cu_seqlens, max_seqlen) + + hidden_states = self.ln_post(hidden_states) + hidden_states, _ = self.proj1(hidden_states) + hidden_states = self.act(hidden_states) + hidden_states, _ = self.proj2(hidden_states) + return hidden_states + + @staticmethod + def map_weight_name(name: str) -> tuple[str, str | None]: + """Map an HF audio-tower name to ``(TokenSpeed name, shard id)``.""" + marker = "audio_tower." + if marker in name: + name = name.split(marker, 1)[1] + name = name.replace("self_attn.out_proj.", "self_attn.proj.") + for checkpoint_name, shard_id in ( + ("q_proj", "q"), + ("k_proj", "k"), + ("v_proj", "v"), + ): + needle = f"self_attn.{checkpoint_name}." + if needle in name: + return name.replace(needle, "self_attn.qkv_proj."), shard_id + return name, None + + def load_weight(self, name: str, loaded_weight: torch.Tensor) -> str | None: + """Load one audio-tower checkpoint tensor. + + ``name`` may be relative to the tower or retain a + ``thinker.audio_tower``/``audio_tower`` prefix. Q/K/V checkpoint + projections are fused into TokenSpeed's ``qkv_proj`` parameter and the + HF ``out_proj`` name is mapped to ``MultimodalEncoderAttention.proj``. + """ + name, shard_id = self.map_weight_name(name) + params = dict(self.named_parameters(remove_duplicate=False)) + if name not in params: + return None + param = params[name] + if shard_id is not None: + param.weight_loader(param, loaded_weight, shard_id) + else: + loader = getattr(param, "weight_loader", default_weight_loader) + loader(param, loaded_weight) + return name + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loaded: set[str] = set() + for name, tensor in weights: + loaded_name = self.load_weight(name, tensor) + if loaded_name is not None: + loaded.add(loaded_name) + return loaded + + +__all__ = [ + "Qwen3AudioEncoder", + "Qwen3AudioEncoderLayer", + "pack_qwen3_audio_features", + "qwen3_audio_output_lengths", +] diff --git a/python/tokenspeed/runtime/models/qwen3_moe.py b/python/tokenspeed/runtime/models/qwen3_moe.py new file mode 100644 index 0000000..3cdfb4d --- /dev/null +++ b/python/tokenspeed/runtime/models/qwen3_moe.py @@ -0,0 +1,308 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Inference-only Qwen3 MoE model compatible with HuggingFace weights.""" + +from __future__ import annotations + +from collections.abc import Iterable + +import torch + +from tokenspeed.runtime.configs.qwen3_moe_config import Qwen3MoeConfig +from tokenspeed.runtime.distributed.comm_manager import CommManager +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.layers.moe import ( + ExpertCheckpointSchema, + build_moe_checkpoint_loader, +) +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.layers.utils import get_layer_id +from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader +from tokenspeed.runtime.models.qwen3 import ( + Qwen3DecoderLayer, + Qwen3ForCausalLM, + Qwen3MLP, + Qwen3Model, +) +from tokenspeed.runtime.models.qwen3_5_moe import ( + Qwen3_5MoeMLP, + Qwen3_5MoeSparseMoeBlock, + _is_moe_layer, +) +from tokenspeed.runtime.utils import add_prefix + + +class Qwen3MoeDecoderLayer(Qwen3DecoderLayer): + def __init__( + self, + config: Qwen3MoeConfig, + mapping: Mapping, + layer_id: int = 0, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__( + config=config, + mapping=mapping, + layer_id=layer_id, + quant_config=quant_config, + prefix=prefix, + ) + if _is_moe_layer(layer_id, config): + self.mlp = Qwen3_5MoeSparseMoeBlock( + config=config, + mapping=self.mapping, + quant_config=quant_config, + layer_index=layer_id, + prefix=add_prefix("mlp", prefix), + ) + elif isinstance(self.mlp, Qwen3MLP): + self.mlp = Qwen3_5MoeMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + mapping=self.mapping, + quant_config=quant_config, + reduce_results=False, + prefix=add_prefix("mlp", prefix), + ) + is_moe = isinstance(self.mlp, Qwen3_5MoeSparseMoeBlock) + self.comm_manager = CommManager( + mapping=self.mapping, + layer_id=layer_id, + is_moe=is_moe, + prev_is_moe=_is_moe_layer(layer_id - 1, config), + input_layernorm=self.input_layernorm, + post_attn_layernorm=self.post_attention_layernorm, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + residual: torch.Tensor | None, + cos_sin: tuple[torch.Tensor, torch.Tensor] | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if not ctx.forward_mode.is_idle(): + hidden_states, residual = self.comm_manager.input_reduce_norm( + hidden_states, residual + ) + hidden_states = self.comm_manager.pre_attn_comm(hidden_states, ctx) + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ctx=ctx, + out_cache_loc=out_cache_loc, + cos_sin=cos_sin, + ) + hidden_states, residual = self.comm_manager.post_attn_reduce_norm( + hidden_states, residual, ctx + ) + + hidden_states, residual = self.forward_mlp(hidden_states, residual, ctx) + return hidden_states, residual + + def forward_mlp( + self, + hidden_states: torch.Tensor, + residual: torch.Tensor, + ctx: ForwardContext, + ) -> tuple[torch.Tensor, torch.Tensor]: + if isinstance(self.mlp, Qwen3_5MoeSparseMoeBlock): + num_global_tokens, max_num_tokens_per_gpu = ( + self.mlp.comm_manager.get_num_tokens(ctx) + ) + hidden_states = self.mlp( + hidden_states, + num_global_tokens, + max_num_tokens_per_gpu, + ctx, + ) + return hidden_states, residual + + hidden_states = self.comm_manager.pre_mlp_comm(hidden_states, ctx) + hidden_states = self.mlp(hidden_states) + hidden_states, residual = self.comm_manager.post_mlp_fused( + hidden_states, residual, ctx + ) + return hidden_states, residual + + +class Qwen3MoeModel(Qwen3Model): + def __init__( + self, + config: Qwen3MoeConfig, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__( + config=config, + mapping=mapping, + quant_config=quant_config, + prefix=prefix, + decoder_layer_type=Qwen3MoeDecoderLayer, + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + input_embeds: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + if input_embeds is None: + hidden_states = self.embed_tokens(input_ids) + else: + hidden_states = input_embeds + residual = None + + for layer in self.layers: + hidden_states, residual = layer( + positions, + hidden_states, + ctx, + out_cache_loc, + residual, + cos_sin=None, + ) + + if not ctx.forward_mode.is_idle(): + hidden_states, _ = layer.comm_manager.final_norm( + hidden_states, residual, ctx, self.norm + ) + return hidden_states, None + + +class Qwen3MoeForCausalLM(Qwen3ForCausalLM): + model_cls = Qwen3MoeModel + + default_bitsandbytes_target_modules = [ + ".gate_proj.", + ".down_proj.", + ".up_proj.", + ".q_proj.", + ".k_proj.", + ".v_proj.", + ".o_proj.", + ] + bitsandbytes_stacked_params_mapping = { + "q_proj": ("qkv_proj", 0), + "k_proj": ("qkv_proj", 1), + "v_proj": ("qkv_proj", 2), + "gate_proj": ("gate_up_proj", 0), + "up_proj": ("gate_up_proj", 1), + } + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): + stacked_params_mapping = [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + ignore_suffixes = ( + ".bias", + "_bias", + ".k_scale", + "_k_scale", + ".v_scale", + "_v_scale", + ".weight_scale", + "_weight_scale", + ".input_scale", + "_input_scale", + ) + + params_dict = dict(self.named_parameters(remove_duplicate=False)) + moe_loader = build_moe_checkpoint_loader( + params_dict=params_dict, + expert_schema=ExpertCheckpointSchema( + gate_proj_name="gate_proj", + down_proj_name="down_proj", + up_proj_name="up_proj", + ), + fused_schema=ExpertCheckpointSchema( + gate_up_fused_name="gate_up_proj", + down_proj_name="down_proj", + ), + num_experts=self.config.num_experts, + ep_rank=self.mapping.moe.ep_rank, + ep_size=self.mapping.moe.ep_size, + ) + + for name, loaded_weight in weights: + if "Embedding" in self.config.name_or_path: + name = add_prefix(name, "model") + layer_id = get_layer_id(name) + if ( + layer_id is not None + and hasattr(self.model, "start_layer") + and ( + layer_id < self.model.start_layer + or layer_id >= self.model.end_layer + ) + ): + continue + if "rotary_emb.inv_freq" in name or "projector" in name: + continue + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + continue + if self.config.tie_word_embeddings and "lm_head.weight" in name: + continue + if name.startswith("model.vision_tower") and name not in params_dict: + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if "mlp.experts" in name: + continue + name = name.replace(weight_name, param_name) + if name.endswith(ignore_suffixes) and name not in params_dict: + continue + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + if name.endswith((".bias", "_bias")) and name not in params_dict: + continue + if moe_loader.matches(name): + moe_loader.load(name, loaded_weight) + continue + if name.endswith(ignore_suffixes) and name not in params_dict: + continue + if name not in params_dict: + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + +EntryClass = Qwen3MoeForCausalLM diff --git a/python/tokenspeed/runtime/models/qwen3_omni.py b/python/tokenspeed/runtime/models/qwen3_omni.py new file mode 100644 index 0000000..9e5fbcf --- /dev/null +++ b/python/tokenspeed/runtime/models/qwen3_omni.py @@ -0,0 +1,485 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Inference-only Qwen3-Omni thinker (text output only). + +The thinker composes TokenSpeed's shared Qwen3 MoE language model, Qwen3 vision +tower, and Qwen3-Omni audio tower. Talker/Code2Wav weights are deliberately +ignored: this entry point implements ``return_audio=False`` and independent +image/audio/video inputs (``use_audio_in_video=False``). +""" + +from __future__ import annotations + +import logging +from collections.abc import Iterable + +import torch + +from tokenspeed.runtime.configs.qwen3_vision_config import Qwen3VLVisionConfig +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.execution.context import ForwardContext +from tokenspeed.runtime.layers.logits_processor import LogitsMetadata +from tokenspeed.runtime.layers.moe import ( + ExpertCheckpointSchema, + build_moe_checkpoint_loader, +) +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.layers.utils import get_layer_id +from tokenspeed.runtime.model_loader.weight_utils import default_weight_loader +from tokenspeed.runtime.models.qwen3_audio import Qwen3AudioEncoder +from tokenspeed.runtime.models.qwen3_moe import Qwen3MoeForCausalLM, Qwen3MoeModel +from tokenspeed.runtime.models.qwen3_vision import Qwen3VLMoeVisionModel +from tokenspeed.runtime.multimodal.embedder import ( + EncoderSpec, + MultimodalEmbedder, + pad_input_tokens, +) +from tokenspeed.runtime.multimodal.encoder_cudagraph import ( + EncoderCudaGraphWrapper, + VisionEncoderCudaGraphAdapter, +) +from tokenspeed.runtime.multimodal.inputs import ( + Modality, + MultimodalDataItem, + MultimodalInputs, +) +from tokenspeed.runtime.utils.env import envs + +logger = logging.getLogger(__name__) + + +def _get_thinker_config(config): + return getattr(config, "thinker_config", config) + + +def _shared_vision_config(vision_config) -> Qwen3VLVisionConfig: + if not getattr(vision_config, "apply_vit_abs_pos_embed", True): + raise ValueError("Qwen3-Omni without absolute vision positions is unsupported") + values = ( + vision_config.to_dict() + if hasattr(vision_config, "to_dict") + else dict(vars(vision_config)) + ) + patch_size = int(values.get("patch_size", 16)) + image_size = values.get("image_size") + if image_size is not None: + image_size = int(image_size) + if image_size % patch_size: + raise ValueError("Qwen3-Omni image_size must be divisible by patch_size") + expected_positions = (image_size // patch_size) ** 2 + configured_positions = values.get("num_position_embeddings") + if ( + configured_positions is not None + and int(configured_positions) != expected_positions + ): + raise ValueError( + "Qwen3-Omni image_size/patch_size disagrees with " + "num_position_embeddings" + ) + values["num_position_embeddings"] = expected_positions + values["deepstack_visual_indexes"] = ( + values.get("deepstack_visual_indexes", [8, 16, 24]) or [] + ) + return Qwen3VLVisionConfig(**values) + + +class Qwen3OmniMoeTextModel(Qwen3MoeModel): + """Qwen3 MoE decoder with Omni visual deepstack injection.""" + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + ctx: ForwardContext, + out_cache_loc: torch.Tensor, + input_embeds: torch.Tensor | None = None, + input_deepstack_embeds: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + hidden_states = ( + self.embed_tokens(input_ids) if input_embeds is None else input_embeds + ) + residual = None + hidden_size = self.config.hidden_size + num_deepstack = ( + input_deepstack_embeds.shape[-1] // hidden_size + if input_deepstack_embeds is not None + else 0 + ) + + for layer_idx, layer in enumerate(self.layers): + hidden_states, residual = layer( + positions, + hidden_states, + ctx, + out_cache_loc, + residual, + cos_sin=None, + ) + if layer_idx < num_deepstack and input_deepstack_embeds.numel() > 0: + start = layer_idx * hidden_size + hidden_states.add_( + input_deepstack_embeds[:, start : start + hidden_size] + ) + + if not ctx.forward_mode.is_idle(): + hidden_states, _ = layer.comm_manager.final_norm( + hidden_states, residual, ctx, self.norm + ) + return hidden_states, None + + +class Qwen3OmniMoeForConditionalGeneration(Qwen3MoeForCausalLM): + """Qwen3-Omni thinker for text-only generation.""" + + model_cls = Qwen3OmniMoeTextModel + + def __init__( + self, + config, + mapping: Mapping, + quant_config: QuantizationConfig | None = None, + is_multimodal_active: bool = True, + mm_attention_backend: str | None = None, + ) -> None: + self.omni_config = config + self.thinker_config = _get_thinker_config(config) + text_config = self.thinker_config.text_config + super().__init__(config=text_config, mapping=mapping, quant_config=quant_config) + + self.is_multimodal_active = is_multimodal_active + self.multimodal_embedder = ( + MultimodalEmbedder() if is_multimodal_active else None + ) + if not is_multimodal_active: + self.visual = None + self.audio_tower = None + self.deepstack_visual_indexes = [] + self.num_deepstack_embeddings = 0 + self.image_encoder = None + self.video_encoder = None + self.audio_encoder = None + return + + vision_config = _shared_vision_config(self.thinker_config.vision_config) + if vision_config.out_hidden_size != text_config.hidden_size: + raise ValueError( + "Qwen3-Omni vision output size must match text hidden size" + ) + if self.thinker_config.audio_config.output_dim != text_config.hidden_size: + raise ValueError("Qwen3-Omni audio output size must match text hidden size") + + self.visual = Qwen3VLMoeVisionModel( + vision_config, + mapping=mapping, + quant_config=None, + norm_eps=getattr(text_config, "rms_norm_eps", 1e-6), + prefix="visual", + mm_attention_backend=mm_attention_backend, + ) + self.audio_tower = Qwen3AudioEncoder( + self.thinker_config.audio_config, + mapping=mapping, + quant_config=None, + prefix="audio_tower", + mm_attention_backend=mm_attention_backend, + ) + self.deepstack_visual_indexes = self.visual.deepstack_visual_indexes + self.num_deepstack_embeddings = len(self.deepstack_visual_indexes) + + # Encoder callables can be replaced by CUDA graph wrappers at runtime. + self.image_encoder = self.get_image_feature + self.video_encoder = self.get_video_feature + self.audio_encoder = self.audio_tower.encode + + def separate_deepstack_embeds( + self, embedding: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + expected_parts = 1 + self.num_deepstack_embeddings + if embedding.shape[-1] != self.config.hidden_size * expected_parts: + raise ValueError( + f"vision embedding width {embedding.shape[-1]} does not match " + f"{expected_parts} x text hidden size {self.config.hidden_size}" + ) + split = self.config.hidden_size + return embedding[:, :split], embedding[:, split:] + + def pad_input_ids( + self, input_ids: list[int], mm_inputs: MultimodalInputs + ) -> list[int]: + return pad_input_tokens(input_ids, mm_inputs) + + def pre_encode( + self, items: list[MultimodalDataItem] + ) -> tuple[torch.Tensor, torch.Tensor]: + pixel_values = torch.cat([item.feature for item in items], dim=0).type( + self.visual.dtype + ) + grid = torch.cat( + [ + ( + item.video_grid_thw + if item.modality == Modality.VIDEO + else item.image_grid_thw + ) + for item in items + ], + dim=0, + ) + if pixel_values.dim() != 2 or grid.dim() != 2: + raise ValueError("Qwen3-Omni vision features require 2-D patches and grids") + return self.visual.prepare_patch_embed(pixel_values, grid), grid + + def post_encode( + self, encoder_outs: list[torch.Tensor], grid: torch.Tensor + ) -> torch.Tensor: + del grid + return torch.cat(encoder_outs, dim=0) + + def get_image_feature(self, items: list[MultimodalDataItem]) -> torch.Tensor: + tokens, grid = self.pre_encode(items) + output = self.visual.forward_blocks(tokens, self.visual.prepare_metadata(grid)) + return self.post_encode([output], grid) + + def get_video_feature(self, items: list[MultimodalDataItem]) -> torch.Tensor: + tokens, grid = self.pre_encode(items) + output = self.visual.forward_blocks(tokens, self.visual.prepare_metadata(grid)) + return self.post_encode([output], grid) + + def _build_encoder_cudagraph_wrapper( + self, + mapping: Mapping, + *, + max_metadata_sequences_per_batch: int | None = None, + metadata_sequence_budget_from_encoder_output_budget: bool = False, + ) -> EncoderCudaGraphWrapper: + adapter = VisionEncoderCudaGraphAdapter( + tower=self.visual, + pre_encode=self.pre_encode, + post_encode=self.post_encode, + out_div=self.visual.spatial_merge_size**2, + merge=self.visual.spatial_merge_size, + input_feature_shape=(1, self.visual.hidden_size), + modality_name="vision", + capture_tp_size=mapping.vision.tp_size, + capture_tp_group=mapping.vision.tp_group, + ) + return EncoderCudaGraphWrapper( + adapter=adapter, + budget_range=(64, 4096), + max_metadata_sequences_per_batch=max_metadata_sequences_per_batch, + metadata_sequence_budget_from_encoder_output_budget=( + metadata_sequence_budget_from_encoder_output_budget + ), + ) + + def make_encoder_cudagraph_wrappers(self, mapping: Mapping) -> dict: + max_video_sequences = ( + envs.TOKENSPEED_MM_VIDEO_ENCODER_CUDA_GRAPH_MAX_SEQUENCES_PER_BATCH.get() + ) + if max_video_sequences is not None: + max_video_sequences = max(1, max_video_sequences) + shared = self._build_encoder_cudagraph_wrapper( + mapping, + max_metadata_sequences_per_batch=max_video_sequences, + metadata_sequence_budget_from_encoder_output_budget=( + max_video_sequences is None + ), + ) + return {"image_encoder": shared, "video_encoder": shared} + + @torch.no_grad() + def forward( + self, + ctx: ForwardContext, + input_ids: torch.Tensor, + positions: torch.Tensor, + out_cache_loc: torch.Tensor, + **kwargs, + ) -> torch.Tensor: + multimodal_context = kwargs.pop("multimodal_context", None) + if ( + not self.is_multimodal_active + and multimodal_context is not None + and multimodal_context.has_extend_inputs() + ): + raise RuntimeError( + "Qwen3-Omni received multimodal inputs while its encoders are disabled" + ) + if ( + multimodal_context is None + or not multimodal_context.has_extend_inputs() + or ctx.forward_mode.is_decode_or_idle() + ): + return super().forward(ctx, input_ids, positions, out_cache_loc, **kwargs) + + input_embeds, model_kwargs = self.multimodal_embedder.apply( + input_ids=input_ids, + text_embedding=self.model.embed_tokens, + ctx=multimodal_context, + encoders={ + Modality.IMAGE: EncoderSpec(self.image_encoder, deepstack=True), + Modality.VIDEO: EncoderSpec(self.video_encoder, deepstack=True), + Modality.AUDIO: EncoderSpec(self.audio_encoder), + }, + multimodal_model=self, + is_decode_or_idle=ctx.forward_mode.is_decode_or_idle(), + ) + hidden_states, aux_hidden_states = self.model( + input_ids, + positions, + ctx, + out_cache_loc, + input_embeds=input_embeds, + **model_kwargs, + ) + return self.logits_processor( + input_ids, + hidden_states, + self.lm_head, + LogitsMetadata.from_forward_context(ctx), + aux_hidden_states, + ) + + @staticmethod + def _map_visual_weight(name: str) -> str: + name = name.replace("attn.qkv.", "attn.qkv_proj.") + name = name.replace("merger_list.", "deepstack_merger_list.") + name = name.replace(".ln_q.", ".norm.") + name = name.replace(".mlp.0.", ".linear_fc1.") + name = name.replace(".mlp.2.", ".linear_fc2.") + return name + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + ignore_suffixes = ( + ".bias", + "_bias", + ".k_scale", + "_k_scale", + ".v_scale", + "_v_scale", + ".weight_scale", + "_weight_scale", + ".input_scale", + "_input_scale", + ) + params = dict(self.named_parameters(remove_duplicate=False)) + moe_loader = build_moe_checkpoint_loader( + params_dict=params, + expert_schema=ExpertCheckpointSchema( + gate_proj_name="gate_proj", + down_proj_name="down_proj", + up_proj_name="up_proj", + ), + fused_schema=ExpertCheckpointSchema( + gate_up_fused_name="gate_up_proj", + down_proj_name="down_proj", + ), + num_experts=self.config.num_experts, + ep_rank=self.mapping.moe.ep_rank, + ep_size=self.mapping.moe.ep_size, + ) + loaded = set() + + for original_name, loaded_weight in weights: + if original_name.startswith(("talker.", "code2wav.")): + continue + name = ( + original_name[len("thinker.") :] + if original_name.startswith("thinker.") + else original_name + ) + + if name.startswith("visual."): + if not self.is_multimodal_active: + continue + name = self._map_visual_weight(name) + if name not in params: + logger.warning("Parameter %s not found in Qwen3-Omni", name) + continue + param = params[name] + loader = getattr(param, "weight_loader", default_weight_loader) + loader(param, loaded_weight) + loaded.add(name) + continue + + if name.startswith("audio_tower."): + if not self.is_multimodal_active: + continue + loaded_name = self.audio_tower.load_weight(name, loaded_weight) + if loaded_name is None: + logger.warning("Parameter %s not found in Qwen3-Omni", name) + else: + loaded.add(f"audio_tower.{loaded_name}") + continue + + layer_id = get_layer_id(name) + if ( + layer_id is not None + and hasattr(self.model, "start_layer") + and not self.model.start_layer <= layer_id < self.model.end_layer + ): + continue + if "rotary_emb" in name: + continue + if self.config.tie_word_embeddings and name == "lm_head.weight": + continue + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name or "mlp.experts" in name: + continue + mapped_name = name.replace(weight_name, param_name) + if mapped_name.endswith(ignore_suffixes) and mapped_name not in params: + break + if mapped_name not in params: + break + param = params[mapped_name] + param.weight_loader(param, loaded_weight, shard_id) + loaded.add(mapped_name) + break + else: + if name.endswith((".bias", "_bias")) and name not in params: + continue + if moe_loader.matches(name): + loaded.add(moe_loader.load(name, loaded_weight)) + continue + if moe_loader.is_expert_checkpoint_weight(name): + continue + if name.endswith(ignore_suffixes) and name not in params: + continue + if name not in params: + logger.warning("Parameter %s not found in Qwen3-Omni", name) + continue + param = params[name] + loader = getattr(param, "weight_loader", default_weight_loader) + loader(param, loaded_weight) + loaded.add(name) + + return loaded + + +EntryClass = [Qwen3OmniMoeForConditionalGeneration] diff --git a/python/tokenspeed/runtime/models/qwen3_vision.py b/python/tokenspeed/runtime/models/qwen3_vision.py new file mode 100644 index 0000000..0b79f1a --- /dev/null +++ b/python/tokenspeed/runtime/models/qwen3_vision.py @@ -0,0 +1,650 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Qwen3 visual tower base reused by the Qwen3.5 target models.""" + +from __future__ import annotations + +from collections.abc import Callable +from functools import lru_cache, partial + +import numpy as np +import torch +import torch.nn as nn +from einops import rearrange +from transformers.activations import ACT2FN + +from tokenspeed.runtime.configs.qwen3_vision_config import Qwen3VLVisionConfig +from tokenspeed.runtime.distributed.mapping import Mapping +from tokenspeed.runtime.layers.attention.mm_encoder_attention import ( + VIT_CUDNN_BATCH_BUCKETS, + VIT_CUDNN_SEQLEN_BUCKETS, + VIT_CUDNN_WORKSPACE_BYTES, + VisionAttention, + round_up_to_bucket, +) +from tokenspeed.runtime.layers.conv import Conv3dLayer +from tokenspeed.runtime.layers.linear import ColumnParallelLinear, RowParallelLinear +from tokenspeed.runtime.layers.quantization.base_config import QuantizationConfig +from tokenspeed.runtime.layers.rotary_embedding import get_rope +from tokenspeed.runtime.layers.vocab_parallel_embedding import VocabParallelEmbedding +from tokenspeed.runtime.utils import add_prefix + + +@lru_cache(maxsize=1024) +def _rot_pos_ids(h: int, w: int, spatial_merge_size: int) -> torch.Tensor: + if isinstance(h, torch.Tensor): + h = int(h.item()) + if isinstance(w, torch.Tensor): + w = int(w.item()) + if isinstance(spatial_merge_size, torch.Tensor): + spatial_merge_size = int(spatial_merge_size.item()) + hpos_ids = np.broadcast_to(np.arange(h).reshape(h, 1), (h, w)) + h_div = h // spatial_merge_size + w_div = w // spatial_merge_size + hpos_ids = hpos_ids.reshape(h_div, spatial_merge_size, w_div, spatial_merge_size) + hpos_ids = hpos_ids.transpose(0, 2, 1, 3).flatten() + + wpos_ids = np.broadcast_to(np.arange(w).reshape(1, w), (h, w)) + wpos_ids = wpos_ids.reshape(h_div, spatial_merge_size, w_div, spatial_merge_size) + wpos_ids = wpos_ids.transpose(0, 2, 1, 3).flatten() + return torch.from_numpy(np.stack([hpos_ids, wpos_ids], axis=-1)) + + +class Qwen3VLVisionMLP(nn.Module): + + def __init__( + self, + in_features: int, + hidden_features: int, + mapping: Mapping, + bias: bool = True, + hidden_act="silu", + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + vision = mapping.vision + self.linear_fc1 = ColumnParallelLinear( + in_features, + hidden_features, + bias=bias, + quant_config=quant_config, + prefix=add_prefix("linear_fc1", prefix), + tp_size=vision.tp_size, + tp_rank=vision.tp_rank, + tp_group=vision.tp_group, + ) + self.linear_fc2 = RowParallelLinear( + hidden_features, + in_features, + bias=bias, + quant_config=quant_config, + prefix=add_prefix("linear_fc2", prefix), + tp_size=vision.tp_size, + tp_rank=vision.tp_rank, + tp_group=vision.tp_group, + reduce_results=True, + ) + self.act = ACT2FN[hidden_act] + + def forward(self, x: torch.Tensor): + x_fc1, _ = self.linear_fc1(x) + x_act = self.act(x_fc1) + mlp_output, _ = self.linear_fc2(x_act) + return mlp_output + + +class Qwen3VLVisionPatchEmbed(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.patch_size = config.patch_size + self.temporal_patch_size = config.temporal_patch_size + self.in_channels = config.in_channels + self.embed_dim = config.hidden_size + + kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size] + self.proj = Conv3dLayer( + self.in_channels, + self.embed_dim, + kernel_size=kernel_size, + stride=kernel_size, + bias=True, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + target_dtype = self.proj.weight.dtype + hidden_states = hidden_states.view( + -1, + self.in_channels, + self.temporal_patch_size, + self.patch_size, + self.patch_size, + ) + hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view( + -1, self.embed_dim + ) + return hidden_states + + +class Qwen3VLVisionBlock(nn.Module): + + def __init__( + self, + dim: int, + num_heads: int, + intermediate_dim: int, + mapping: Mapping, + head_size: int | None = None, + hidden_act="silu", + norm_layer: Callable[[int], nn.Module] | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + workspace_buffer: torch.Tensor | None = None, + mm_attention_backend: str | None = None, + ) -> None: + super().__init__() + if norm_layer is None: + norm_layer = partial(nn.LayerNorm, eps=1e-6) + self.norm1 = norm_layer(dim) + self.norm2 = norm_layer(dim) + + self.attn = VisionAttention( + embed_dim=dim, + num_heads=num_heads, + head_size=head_size, + proj_bias=True, + quant_config=quant_config, + prefix=add_prefix("attn", prefix), + workspace_buffer=workspace_buffer, + mapping=mapping, + mm_attention_backend=mm_attention_backend, + ) + self.mlp = Qwen3VLVisionMLP( + dim, + intermediate_dim, + hidden_act=hidden_act, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + mapping=mapping, + ) + + def forward( + self, + x: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb_cos: torch.Tensor, + rotary_pos_emb_sin: torch.Tensor, + max_seqlen: int | None = None, + sequence_lengths: torch.Tensor | None = None, + ) -> torch.Tensor: + hidden_states = self.norm1(x) + hidden_states = rearrange(hidden_states, "s b ... -> b s ...") + attn = self.attn( + hidden_states, + cu_seqlens=cu_seqlens, + rotary_pos_emb_cos=rotary_pos_emb_cos, + rotary_pos_emb_sin=rotary_pos_emb_sin, + max_seqlen=max_seqlen, + sequence_lengths=sequence_lengths, + ) + attn = rearrange(attn, "b s ... -> s b ...") + x += attn + norm2 = self.norm2(x) + mlp = self.mlp(norm2) + x += mlp + return x + + +class Qwen3VLMoeVisionPatchMerger(nn.Module): + + def __init__( + self, + dim: int, + context_dim: int, + padded_context_dim: int, + mapping: Mapping, + norm_layer: Callable[[int], nn.Module] | None = None, + spatial_merge_size: int = 2, + use_postshuffle_norm: bool = False, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = context_dim * (spatial_merge_size**2) + self.padded_context_dim = padded_context_dim * (spatial_merge_size**2) + + self.use_postshuffle_norm = use_postshuffle_norm + + if norm_layer is None: + norm_layer = partial(nn.LayerNorm, eps=1e-6) + self.norm = norm_layer( + self.hidden_size if use_postshuffle_norm else context_dim + ) + vision = mapping.vision + self.linear_fc1 = ColumnParallelLinear( + self.hidden_size, + self.padded_context_dim, + bias=True, + quant_config=quant_config, + prefix=add_prefix("linear_fc1", prefix), + tp_size=vision.tp_size, + tp_rank=vision.tp_rank, + tp_group=vision.tp_group, + ) + self.act_fn = nn.GELU() + self.linear_fc2 = RowParallelLinear( + self.padded_context_dim, + dim, + bias=True, + quant_config=quant_config, + prefix=add_prefix("linear_fc2", prefix), + tp_size=vision.tp_size, + tp_rank=vision.tp_rank, + tp_group=vision.tp_group, + reduce_results=True, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.use_postshuffle_norm: + x = self.norm(x.view(-1, self.hidden_size)) + else: + x = self.norm(x).view(-1, self.hidden_size) + + x_parallel, _ = self.linear_fc1(x) + x_parallel = self.act_fn(x_parallel) + out, _ = self.linear_fc2(x_parallel) + return out + + +class Qwen3VLMoeVisionModel(nn.Module): + + def __init__( + self, + vision_config: Qwen3VLVisionConfig, + mapping: Mapping, + norm_eps: float = 1e-6, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + mm_attention_backend: str | None = None, + ) -> None: + super().__init__() + vision = mapping.vision + self.mm_attention_backend = mm_attention_backend + self.hidden_size = vision_config.hidden_size + self.num_heads = vision_config.num_heads + self.num_position_embeddings = vision_config.num_position_embeddings + self.num_grid_per_side = int(self.num_position_embeddings**0.5) + self.spatial_merge_size = vision_config.spatial_merge_size + # layer indices whose outputs feed the deepstack mergers + self.deepstack_visual_indexes = vision_config.deepstack_visual_indexes + self.patch_embed = Qwen3VLVisionPatchEmbed(config=vision_config) + self.pos_embed = VocabParallelEmbedding( + self.num_position_embeddings, + self.hidden_size, + quant_config=quant_config, + tp_rank=vision.tp_rank, + tp_size=vision.tp_size, + tp_group=vision.tp_group, + prefix=add_prefix("pos_embed", prefix), + ) + + norm_layer = partial(nn.LayerNorm, eps=norm_eps) + head_dim = self.hidden_size // self.num_heads + self.rotary_pos_emb = get_rope( + head_size=head_dim, + rotary_dim=head_dim // 2, + max_position=8192, + base=10000.0, + is_neox_style=True, + ) + + workspace_buffer = None + if self.mm_attention_backend == "flashinfer_cudnn": + if torch.cuda.is_available(): + ws_device = torch.device("cuda", torch.cuda.current_device()) + else: + ws_device = self.device + workspace_buffer = torch.empty( + VIT_CUDNN_WORKSPACE_BYTES, + dtype=torch.uint8, + device=ws_device, + ) + + self.blocks = nn.ModuleList( + [ + Qwen3VLVisionBlock( + dim=self.hidden_size, + num_heads=self.num_heads, + intermediate_dim=vision_config.intermediate_size, + head_size=head_dim, + hidden_act=vision_config.hidden_act, + norm_layer=norm_layer, + quant_config=quant_config, + prefix=add_prefix(f"blocks.{layer_idx}", prefix), + workspace_buffer=workspace_buffer, + mapping=mapping, + mm_attention_backend=mm_attention_backend, + ) + for layer_idx in range(vision_config.depth) + ] + ) + self.merger = Qwen3VLMoeVisionPatchMerger( + dim=vision_config.out_hidden_size, + context_dim=self.hidden_size, + padded_context_dim=self.num_heads * head_dim, + norm_layer=norm_layer, + spatial_merge_size=self.spatial_merge_size, + quant_config=quant_config, + prefix=add_prefix("merger", prefix), + mapping=mapping, + ) + + self.deepstack_merger_list = nn.ModuleList( + [ + Qwen3VLMoeVisionPatchMerger( + dim=vision_config.out_hidden_size, + context_dim=self.hidden_size, + padded_context_dim=self.num_heads * head_dim, + spatial_merge_size=self.spatial_merge_size, + use_postshuffle_norm=True, + norm_layer=norm_layer, + quant_config=quant_config, + prefix=add_prefix(f"deepstack_merger_list.{layer_idx}", prefix), + mapping=mapping, + ) + for layer_idx in range(len(self.deepstack_visual_indexes)) + ] + ) + + self.tp_size = vision.tp_size + + @property + def dtype(self) -> torch.dtype: + return self.patch_embed.proj.weight.dtype + + @property + def device(self) -> torch.device: + return self.patch_embed.proj.weight.device + + def rot_pos_emb( + self, grid_thw: list[list[int]] + ) -> tuple[torch.Tensor, torch.Tensor]: + pos_ids = [] + for t, h, w in grid_thw: + base = _rot_pos_ids(h, w, self.spatial_merge_size) + pos_ids.append(base if t == 1 else base.repeat(t, 1)) + + pos_ids = torch.cat(pos_ids, dim=0).to(self.device, non_blocking=True) + max_grid_size = max(max(h, w) for _, h, w in grid_thw) + + cos, sin = self._get_rotary_cos_sin(max_grid_size) + + cos_combined = cos[pos_ids].flatten(1) + sin_combined = sin[pos_ids].flatten(1) + + return cos_combined, sin_combined + + def _get_rotary_cos_sin(self, seqlen: int) -> tuple[torch.Tensor, torch.Tensor]: + cos_sin = self.rotary_pos_emb.cos_sin_cache[:seqlen].to(self.device) + return cos_sin.chunk(2, dim=-1) + + def fast_pos_embed_interpolate_from_list(self, grid_thw): + num_grid_per_side = self.num_grid_per_side + m_size = self.spatial_merge_size + hidden_dim = self.pos_embed.embedding_dim + + outputs = [] + for t, h, w in grid_thw: + h_idxs = torch.linspace( + 0, num_grid_per_side - 1, h, dtype=torch.float32, device=self.device + ) + w_idxs = torch.linspace( + 0, num_grid_per_side - 1, w, dtype=torch.float32, device=self.device + ) + + h_floor = h_idxs.to(torch.long) + w_floor = w_idxs.to(torch.long) + h_ceil = torch.clamp(h_floor + 1, max=num_grid_per_side - 1) + w_ceil = torch.clamp(w_floor + 1, max=num_grid_per_side - 1) + + dh = h_idxs - h_floor + dw = w_idxs - w_floor + + # Create meshgrid view for all h, w vars + dh_grid, dw_grid = torch.meshgrid(dh, dw, indexing="ij") + h_floor_grid, w_floor_grid = torch.meshgrid(h_floor, w_floor, indexing="ij") + h_ceil_grid, w_ceil_grid = torch.meshgrid(h_ceil, w_ceil, indexing="ij") + + # original computation of weights + # w00 = (1 - dh_grid) * (1 - dw_grid) + # w01 = (1 - dh_grid) * dw_grid + # w10 = dh_grid * (1 - dw_grid) + # w11 = dh_grid * dw_grid + # we reuse w11 here to avoid duplicate + # dh_grid * dw_grid computation + w11 = dh_grid * dw_grid + w10 = dh_grid - w11 + w01 = dw_grid - w11 + w00 = 1 - dh_grid - w01 + + h_grid = torch.stack([h_floor_grid, h_floor_grid, h_ceil_grid, h_ceil_grid]) + w_grid = torch.stack([w_floor_grid, w_ceil_grid, w_floor_grid, w_ceil_grid]) + h_grid_idx = h_grid * num_grid_per_side + + indices = (h_grid_idx + w_grid).reshape(4, -1) + weights = torch.stack([w00, w01, w10, w11], dim=0).reshape(4, -1, 1) + weights = weights.to(dtype=self.dtype) + + embeds = self.pos_embed(indices) + embeds *= weights + combined = embeds.sum(dim=0) + + combined = combined.reshape( + h // m_size, m_size, w // m_size, m_size, hidden_dim + ) + combined = combined.permute(0, 2, 1, 3, 4).reshape(1, -1, hidden_dim) + repeated = combined.expand(t, -1, -1).reshape(-1, hidden_dim) + outputs.append(repeated) + + return torch.cat(outputs, dim=0) + + def compute_cudnn_batch_offsets_packed( + self, + token_cu_seqlens: np.ndarray, + *, + elem_per_token: int, + ) -> np.ndarray: + """ + Build packed *element* indptrs for cuDNN prefill. + + Input: + token_cu_seqlens: (B+1,) token indptr + elem_per_token: per-token element width on THIS TP rank + (usually hidden_size / attn_tp_size) + + Output: + packed_offsets: (3 * (B_padded + 1),) int32 + [qk_indptr, v_indptr, o_indptr] concatenated, + each indptr is (B_padded + 1,) in element units. + """ + if token_cu_seqlens.ndim != 1 or token_cu_seqlens.size < 2: + raise ValueError( + "token_cu_seqlens must be a 1D array with at least 2 entries." + ) + B = int(token_cu_seqlens.size - 1) + B_padded = round_up_to_bucket(B, VIT_CUDNN_BATCH_BUCKETS) + + # token indptr -> pad to (B_padded+1,) by appending total_tokens for extra empty sequences + token_indptr = token_cu_seqlens.astype(np.int64, copy=False) # (B+1,) + if B_padded != B: + pad = np.full((B_padded - B,), token_indptr[-1], dtype=token_indptr.dtype) + token_indptr = np.concatenate([token_indptr, pad], axis=0) # (B_padded+1,) + + # convert token indptr -> element indptr + elem_indptr = (token_indptr * int(elem_per_token)).astype( + np.int32 + ) # (B_padded+1,) + + # q/k/v/o in this vision encoder path share the same indptr + return np.concatenate([elem_indptr, elem_indptr, elem_indptr], axis=0) + + def compute_cudnn_sequence_lengths_padded( + self, + token_cu_seqlens: np.ndarray, + ) -> np.ndarray: + """ + token_cu_seqlens: (B+1,) token indptr + return: (B_padded,) token lengths (padded with 0) + """ + if token_cu_seqlens.ndim != 1 or token_cu_seqlens.size < 2: + raise ValueError( + "token_cu_seqlens must be a 1D array with at least 2 entries." + ) + B = int(token_cu_seqlens.size - 1) + + seq_lens = (token_cu_seqlens[1:] - token_cu_seqlens[:-1]).astype( + np.int32 + ) # (B,) + + B_padded = round_up_to_bucket(B, VIT_CUDNN_BATCH_BUCKETS) + if B_padded != B: + pad = np.zeros((B_padded - B,), dtype=np.int32) + seq_lens = np.concatenate([seq_lens, pad], axis=0) # (B_padded,) + return seq_lens + + def prepare_patch_embed( + self, x: torch.Tensor, grid_thw: torch.Tensor | list + ) -> torch.Tensor: + """Eager patch-embed (runs before the captured region): Conv patch embed + + interpolated position embedding + the ``[s, 1, h]`` reshape the block + loop expects. + + Kept eager (outside the capture-safe region) -- the interpolation does + host/numpy work. + """ + x = x.to(device=self.device, dtype=self.dtype) + x = self.patch_embed(x) + grid_thw_list = grid_thw if isinstance(grid_thw, list) else grid_thw.tolist() + x = x + self.fast_pos_embed_interpolate_from_list(grid_thw_list) + return x.unsqueeze(1) + + def prepare_metadata(self, grid_thw: torch.Tensor | list) -> dict: + """Eager metadata pass: rotary embeddings, cu_seqlens, sequence lengths, + and ``max_seqlen`` as a Python int. + + Everything here involves a host sync or a data-dependent shape, so it + lives outside the capture-safe block loop. ``max_seqlen`` is + materialized as a plain int (CPU/numpy, no GPU sync) so the captured + block loop never hits the attention backend's ``.item()`` fallback. + """ + if isinstance(grid_thw, list): + grid_thw_list = grid_thw + grid_thw_np = np.array(grid_thw, dtype=np.int32) + else: + grid_thw_list = grid_thw.tolist() + grid_thw_np = grid_thw.cpu().numpy() + + rotary_pos_emb_cos, rotary_pos_emb_sin = self.rot_pos_emb(grid_thw_list) + + # ---- build token indptr (B+1,) ---- + token_cu_seqlens = np.concatenate( + [ + np.zeros(1, dtype=np.int32), + np.repeat( + grid_thw_np[:, 1] * grid_thw_np[:, 2], grid_thw_np[:, 0] + ).cumsum(axis=0, dtype=np.int32), + ] + ) + real_seq_lens = token_cu_seqlens[1:] - token_cu_seqlens[:-1] + real_max_seqlen = int(real_seq_lens.max()) if real_seq_lens.size > 0 else 0 + + if self.mm_attention_backend == "flashinfer_cudnn": + # (B_padded,) token lengths + seq_lens_padded = self.compute_cudnn_sequence_lengths_padded( + token_cu_seqlens + ) + # element-per-token width on this vision TP rank + elem_per_token = ( + self.hidden_size // self.tp_size + ) # == heads_per_rank * head_dim + # (3*(B_padded+1),) packed element indptrs + offsets_packed = self.compute_cudnn_batch_offsets_packed( + token_cu_seqlens, + elem_per_token=elem_per_token, + ) + sequence_lengths = ( + torch.from_numpy(seq_lens_padded) + .to(device=self.device, dtype=torch.int32, non_blocking=True) + .view(-1, 1, 1, 1) + ) # match cuDNN test style + cu_seqlens = torch.from_numpy(offsets_packed).to( + device=self.device, dtype=torch.int32, non_blocking=True + ) + max_seqlen = round_up_to_bucket(real_max_seqlen, VIT_CUDNN_SEQLEN_BUCKETS) + else: + sequence_lengths = None + cu_seqlens = torch.from_numpy(token_cu_seqlens).to( + device=self.device, dtype=torch.int32, non_blocking=True + ) + max_seqlen = real_max_seqlen + + return { + "cu_seqlens": cu_seqlens, + "rotary_pos_emb_cos": rotary_pos_emb_cos, + "rotary_pos_emb_sin": rotary_pos_emb_sin, + "max_seqlen": max_seqlen, + "sequence_lengths": sequence_lengths, + } + + def forward_blocks(self, x: torch.Tensor, metadata: dict) -> torch.Tensor: + """Capture-safe encoder body: block loop + deepstack mergers + merger. + + No host syncs and no data-dependent control flow, so this region is + safe to record into a CUDA graph. ``metadata`` comes from + :meth:`prepare_metadata`; ``x`` from :meth:`prepare_patch_embed`. + """ + cu_seqlens = metadata["cu_seqlens"] + rotary_pos_emb_cos = metadata["rotary_pos_emb_cos"] + rotary_pos_emb_sin = metadata["rotary_pos_emb_sin"] + max_seqlen = metadata["max_seqlen"] + sequence_lengths = metadata["sequence_lengths"] + + deepstack_feature_lists = [] + num_deepstack_captured = 0 + for layer_num, blk in enumerate(self.blocks): + x = blk( + x, + cu_seqlens=cu_seqlens, + rotary_pos_emb_cos=rotary_pos_emb_cos, + rotary_pos_emb_sin=rotary_pos_emb_sin, + max_seqlen=max_seqlen, + sequence_lengths=sequence_lengths, + ) + if layer_num in self.deepstack_visual_indexes: + deepstack_feature = self.deepstack_merger_list[num_deepstack_captured]( + x + ) + deepstack_feature_lists.append(deepstack_feature) + num_deepstack_captured += 1 + x = self.merger(x) + # [seq_len, out_hidden_size * (1 + depth_of_deepstack)] + return torch.cat([x] + deepstack_feature_lists, dim=1) diff --git a/python/tokenspeed/runtime/models/registry.py b/python/tokenspeed/runtime/models/registry.py new file mode 100755 index 0000000..1ed6ce0 --- /dev/null +++ b/python/tokenspeed/runtime/models/registry.py @@ -0,0 +1,123 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Model registry helpers for resolving runtime model entry classes.""" + +import importlib +import pkgutil +from collections.abc import Set +from dataclasses import dataclass, field +from functools import lru_cache + +import torch.nn as nn + +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +@dataclass +class _ModelRegistry: + # Keyed by model_arch + models: dict[str, type[nn.Module] | str] = field(default_factory=dict) + + def get_supported_archs(self) -> Set[str]: + return self.models.keys() + + def _raise_for_unsupported(self, architectures: list[str]): + all_supported_archs = self.get_supported_archs() + + if any(arch in all_supported_archs for arch in architectures): + raise ValueError( + f"Model architectures {architectures} failed " + "to be inspected. Please check the logs for more details." + ) + + raise ValueError( + f"Model architectures {architectures} are not supported for now. " + f"Supported architectures: {all_supported_archs}" + ) + + def _try_load_model_cls(self, model_arch: str) -> type[nn.Module] | None: + if model_arch not in self.models: + return None + + return self.models[model_arch] + + def _normalize_archs( + self, + architectures: str | list[str], + ) -> list[str]: + if isinstance(architectures, str): + architectures = [architectures] + if not architectures: + logger.warning("No model architectures are specified") + + return architectures + + def resolve_model_cls( + self, + architectures: str | list[str], + ) -> tuple[type[nn.Module], str]: + architectures = self._normalize_archs(architectures) + + for arch in architectures: + model_cls = self._try_load_model_cls(arch) + if model_cls is not None: + return (model_cls, arch) + + return self._raise_for_unsupported(architectures) + + +@lru_cache +def import_model_classes(): + """Import model modules and collect their ``EntryClass`` exports.""" + model_arch_name_to_cls = {} + package_name = "tokenspeed.runtime.models" + package = importlib.import_module(package_name) + + for _, name, _ in pkgutil.iter_modules(package.__path__, package_name + "."): + + try: + module = importlib.import_module(name) + except Exception as exc: + raise RuntimeError(f"Failed to import model module {name}.") from exc + if hasattr(module, "EntryClass"): + entry = module.EntryClass + if isinstance( + entry, list + ): # To support multiple model classes in one module + for tmp in entry: + if tmp.__name__ in model_arch_name_to_cls: + raise ValueError( + f"Duplicated model implementation for {tmp.__name__}" + ) + model_arch_name_to_cls[tmp.__name__] = tmp + else: + if entry.__name__ in model_arch_name_to_cls: + raise ValueError( + f"Duplicated model implementation for {entry.__name__}" + ) + model_arch_name_to_cls[entry.__name__] = entry + + return model_arch_name_to_cls + + +ModelRegistry = _ModelRegistry(import_model_classes()) diff --git a/python/tokenspeed/runtime/models/utils.py b/python/tokenspeed/runtime/models/utils.py new file mode 100755 index 0000000..0064e85 --- /dev/null +++ b/python/tokenspeed/runtime/models/utils.py @@ -0,0 +1,100 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Helpers shared across runtime model implementations.""" + +import torch +from tokenspeed_kernel.ops.embedding import FusedSetKVBufferArg + +from tokenspeed.runtime.layers.paged_attention import PagedAttention +from tokenspeed.runtime.utils import print_warning_once + + +def validate_attention_partition( + total_num_heads: int, + total_num_kv_heads: int, + tp_size: int, +) -> None: + if tp_size <= 0: + raise ValueError(f"tp_size must be positive, got {tp_size}.") + if total_num_heads % tp_size != 0: + raise ValueError( + f"num_attention_heads={total_num_heads} must be divisible by tp_size={tp_size}." + ) + if total_num_kv_heads <= 0: + raise ValueError( + f"num_key_value_heads must be positive, got {total_num_kv_heads}." + ) + if total_num_kv_heads >= tp_size: + if total_num_kv_heads % tp_size != 0: + raise ValueError( + f"num_key_value_heads={total_num_kv_heads} must be divisible by tp_size={tp_size}." + ) + elif tp_size % total_num_kv_heads != 0: + raise ValueError( + f"tp_size={tp_size} must be divisible by num_key_value_heads={total_num_kv_heads}." + ) + + +def create_fused_set_kv_buffer_arg( + value: torch.Tensor, + layer: PagedAttention, + out_cache_loc: torch.Tensor, + token_to_kv_pool, +): + """Build fused RoPE+KV write arguments when the fused path is supported.""" + + from tokenspeed.runtime.layers.attention.kv_cache.mla import MLATokenToKVPool + + layer_id = layer.layer_id + + k_buffer = token_to_kv_pool.get_key_buffer(layer_id) + v_buffer = token_to_kv_pool.get_value_buffer(layer_id) + + is_mla = isinstance(token_to_kv_pool, MLATokenToKVPool) + + if is_mla: + kv_lora_rank = token_to_kv_pool.kv_lora_rank + k_buffer = k_buffer[..., kv_lora_rank:].view(k_buffer.shape[0], -1) + v_buffer = v_buffer[..., :kv_lora_rank].view(v_buffer.shape[0], -1) + else: + k_buffer = k_buffer.view(k_buffer.shape[0], -1) + v_buffer = v_buffer.view(v_buffer.shape[0], -1) + + # Non-trivial scales need 1/scale applied before FP8 cast — the fused kernel + # doesn't support this yet, so log a warning and skip the fused path. + k_scale = layer.k_scale + v_scale = layer.v_scale + if (k_scale is not None and k_scale != 1.0) or ( + v_scale is not None and v_scale != 1.0 + ): + print_warning_once( + f"Fused RoPE+KV write disabled: non-trivial k_scale={k_scale} v_scale={v_scale}" + ) + return None + + return FusedSetKVBufferArg( + value=value, + k_buffer=k_buffer, + v_buffer=v_buffer, + k_scale=None, + v_scale=None, + cache_loc=out_cache_loc, + ) diff --git a/python/tokenspeed/runtime/moe/__init__.py b/python/tokenspeed/runtime/moe/__init__.py new file mode 100644 index 0000000..adda4a3 --- /dev/null +++ b/python/tokenspeed/runtime/moe/__init__.py @@ -0,0 +1,21 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Mixture-of-experts runtime helpers and metadata.""" diff --git a/python/tokenspeed/runtime/moe/distribution_recorder.py b/python/tokenspeed/runtime/moe/distribution_recorder.py new file mode 100755 index 0000000..8a17814 --- /dev/null +++ b/python/tokenspeed/runtime/moe/distribution_recorder.py @@ -0,0 +1,868 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the FluentLLM project +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import logging +import time +from abc import ABC +from collections import deque +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Literal + +import einops +import torch +import torch.distributed + +from tokenspeed.runtime.moe.expert_location import ExpertLocationMetadata +from tokenspeed.runtime.utils import Withable +from tokenspeed.runtime.utils.env import envs +from tokenspeed.runtime.utils.server_args import ServerArgs + +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: + 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 with_forward_pass( + self, forward_pass_id: int, metadata: dict[str, Any] | None = None + ): + yield + + def on_select_experts(self, topk_ids: torch.Tensor, num_experts: int | None = None): + 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 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 RuntimeError( + "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._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, metadata: dict[str, Any] | None = None + ): + with self._current_forward_pass_id.with_value(forward_pass_id): + self._on_forward_pass_start(metadata) + try: + yield + finally: + self._on_forward_pass_end(forward_pass_id) + + def _on_forward_pass_start(self, metadata: dict[str, Any] | None = None): + if not self._recording: + return + for gatherer_key, gatherer in self._single_pass_gatherers.items(): + gatherer.reset() + gatherer.on_forward_pass_start(metadata) + + 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, num_experts: int | None = None): + self._on_hook("on_select_experts", topk_ids=topk_ids, num_experts=num_experts) + + 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_hook(self, hook_name: str, **kwargs): + if not (self._recording or torch.cuda.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...") + if self._current_layer_idx.value is not None: + raise RuntimeError(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( + "TokenSpeed 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( + "TokenSpeed 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: ExpertDistributionRecorder | None = ( + _ExpertDistributionRecorderNoop() +) + + +def get_global_expert_distribution_recorder(): + return _global_expert_distribution_recorder + + +# --------------------------------------- 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.enable_deepep_moe and (server_args.deepep_mode == "normal"): + return _DeepepNormalSinglePassGatherer(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, metadata: dict[str, Any] | None = None): + 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 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: dict[str, Any] | None = None + self._topk_ids_of_layer = torch.zeros( + ( + expert_location_metadata.num_layers, + server_args.chunked_prefill_size * 8, + self._TOP_K_NUM, + ), + dtype=torch.int32, + device=server_args.device, + ) + self._misc_objects: list[dict[str, Any]] = [] + + def on_forward_pass_start(self, metadata: dict[str, Any] | None = None): + if self._metadata is not None: + raise RuntimeError("forward pass metadata was not collected before reset.") + self._metadata = dict(metadata or {}) + + 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]): + if not 0 <= layer_idx < self._expert_location_metadata.num_layers: + raise ValueError( + f"layer_idx={layer_idx} is outside [0, {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) + 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="cuda", + ) + + 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, num_experts: int | None = None + ): + topk_ids = topk_ids.flatten() + if num_experts is None: + mask = topk_ids != -1 + else: + mask = (topk_ids != -1) & (topk_ids < num_experts) + 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, + ): + if not isinstance(local_physical_count_of_layer, list): + raise TypeError( + "local_physical_count_of_layer must be a list, got " + f"{type(local_physical_count_of_layer).__name__}." + ) + 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) + + +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 + + global_count = torch.zeros( + (num_layers, num_physical_experts), dtype=dtype, device=device + ) + global_count[ + :, num_local_physical_experts * rank : num_local_physical_experts * (rank + 1) + ] = local_physical_count + return global_count + + +# --------------------------------------- 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: str | None): + 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 _UtilizationRateAccumulator(_Accumulator): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + self._enable = self._server_args.enable_expert_distribution_metrics + + if self._enable: + window_sizes = [10, 100, 1000] + self._history = _DequeCollection(maxlens=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( + "[Expert Balancedness] forward_pass_id=%s current_pass_balancedness=%.03f %s gpu_physical_count_sum=%s", + forward_pass_id, + utilization_rate, + "".join( + f"last_{size}_average_balancedness={value:.03f} " + for size, value in self._history.mean().items() + ), + 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(_UtilizationRateAccumulator): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._records = [] + + 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): + if output_mode != "file": + raise ValueError(f"Unsupported output_mode: {output_mode}") + output = dict( + records=self._records, + # 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(_UtilizationRateAccumulator): + 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.cuda.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, + ) + + 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 _dump_to_file(name, data): + save_dir = Path(envs.TOKENSPEED_EXPERT_DISTRIBUTION_RECORDER_DIR.get()) + path_output = save_dir / name + logger.info("Write expert distribution to %s", 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/tokenspeed/runtime/moe/eplb_algorithms/__init__.py b/python/tokenspeed/runtime/moe/eplb_algorithms/__init__.py new file mode 100755 index 0000000..0bb5b9b --- /dev/null +++ b/python/tokenspeed/runtime/moe/eplb_algorithms/__init__.py @@ -0,0 +1,65 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from enum import Enum, auto + +import torch + +from tokenspeed.runtime.moe.eplb_algorithms import deepseek + + +class EplbAlgorithm(Enum): + deepseek = auto() + deepseek_hierarchical = auto() + + +def rebalance_experts( + tokens_per_expert: torch.Tensor, + num_physical_experts: int, + num_local_physical_experts: int, + num_groups: int | None, + 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, + ) + + raise NotImplementedError + + +def compute_algorithm( + raw_algorithm: str, + num_groups: int | None, + num_nodes: int, +) -> EplbAlgorithm: + if raw_algorithm != "auto": + return EplbAlgorithm[raw_algorithm] + + if (num_groups is not None) and (num_groups % num_nodes == 0): + return EplbAlgorithm.deepseek_hierarchical + else: + return EplbAlgorithm.deepseek diff --git a/python/tokenspeed/runtime/moe/eplb_algorithms/deepseek.py b/python/tokenspeed/runtime/moe/eplb_algorithms/deepseek.py new file mode 100755 index 0000000..a2c20d9 --- /dev/null +++ b/python/tokenspeed/runtime/moe/eplb_algorithms/deepseek.py @@ -0,0 +1,261 @@ +# SPDX-License-Identifier: MIT +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright (c) 2025 DeepSeek +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import torch + + +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 + if num_packs <= 0 or num_groups % num_packs != 0: + raise ValueError( + f"num_groups={num_groups} must be divisible by num_packs={num_packs}." + ) + 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__, + ) + if pack_items[pack] >= groups_per_pack: + raise RuntimeError("balanced_packing selected a full 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 + if num_redundant < 0: + raise ValueError( + f"num_phy={num_phy} must be greater than or equal to num_log={num_log}." + ) + 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 + if num_groups <= 0 or num_logical_experts % num_groups != 0: + raise ValueError( + f"num_logical_experts={num_logical_experts} must be divisible by num_groups={num_groups}." + ) + group_size = num_logical_experts // num_groups + if num_nodes <= 0 or num_groups % num_nodes != 0: + raise ValueError( + f"num_groups={num_groups} must be divisible by num_nodes={num_nodes}." + ) + groups_per_node = num_groups // num_nodes + if num_gpus <= 0 or num_gpus % num_nodes != 0: + raise ValueError( + f"num_gpus={num_gpus} must be divisible by num_nodes={num_nodes}." + ) + if num_physical_experts % num_gpus != 0: + raise ValueError( + f"num_physical_experts={num_physical_experts} must be divisible by num_gpus={num_gpus}." + ) + 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/tokenspeed/runtime/moe/expert_location.py b/python/tokenspeed/runtime/moe/expert_location.py new file mode 100755 index 0000000..abef911 --- /dev/null +++ b/python/tokenspeed/runtime/moe/expert_location.py @@ -0,0 +1,467 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the FluentLLM project +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import json +import logging +import random +from dataclasses import dataclass +from pathlib import Path + +import torch +import torch.distributed +import torch.nn.functional as F + +from tokenspeed.runtime.configs.model_config import ModelConfig +from tokenspeed.runtime.model_loader import get_model_architecture +from tokenspeed.runtime.moe import eplb_algorithms +from tokenspeed.runtime.utils.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_num_valid: torch.Tensor # (layers, num_logical_experts) + # (layers, num_logical_experts) + logical_to_rank_dispatch_physical_map: torch.Tensor | None + + # -------------------------------- 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: + count, remainder = divmod(self.num_physical_experts, self.ep_size) + if remainder != 0: + raise ValueError( + f"num_physical_experts={self.num_physical_experts} must be divisible by ep_size={self.ep_size}." + ) + return count + + @property + def num_logical_experts(self) -> int: + return self.logical_to_all_physical_map.shape[1] + + @property + def ep_size(self): + 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 + ) + if not num_layers_0 == num_layers_1 == num_layers_2: + raise ValueError( + "Expert location maps disagree on layer count: " + f"{num_layers_0}, {num_layers_1}, {num_layers_2}." + ) + if num_logical_experts_0 != num_logical_experts_1: + raise ValueError( + "Expert location maps disagree on logical expert count: " + f"{num_logical_experts_0}, {num_logical_experts_1}." + ) + if num_physical_experts_0 != num_physical_experts_1: + raise ValueError( + "Expert location maps disagree on physical expert count: " + f"{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) + 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) + 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) + 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.mapping.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) + ) + + num_physical_experts = ( + model_config_for_expert_location.num_logical_experts + + server_args.ep_num_redundant_experts + ) + ep_size = server_args.mapping.moe.ep_size + if ep_size <= 0 or num_physical_experts % ep_size != 0: + raise ValueError(f"{num_physical_experts=} {ep_size=}") + 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_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, + ep_rank=torch.distributed.get_rank() % ep_size, + ) + if server_args.ep_dispatch_algorithm == "static" + or server_args.ep_dispatch_algorithm == "static_with_zero_expert" + else None + ), + ) + + # -------------------------------- mutation ------------------------------------ + + def update( + self, + other: "ExpertLocationMetadata", + update_layer_ids: list[int], + ): + for field in [ + "ep_size", + ]: + if getattr(self, field) != getattr(other, field): + raise ValueError( + f"Cannot update ExpertLocationMetadata with different {field}." + ) + + for field in [ + "physical_to_logical_map", + "physical_to_logical_map_cpu", + "logical_to_all_physical_map", + "logical_to_all_physical_map_num_valid", + "logical_to_rank_dispatch_physical_map", + ]: + other_field = getattr(other, field) + self_field = getattr(self, field) + if (other_field is not None) != (self_field is not None): + raise ValueError( + f"Cannot update ExpertLocationMetadata with incompatible {field}." + ) + 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]: + return [ + physical_expert_id + for physical_expert_id in self.logical_to_all_physical_map[ + layer_id, logical_expert_id + ].tolist() + if physical_expert_id != -1 + ] + + +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 + + +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, + ) + + if not torch.all(logical_to_rank_dispatch_physical_map != -1): + raise RuntimeError( + "logical_to_rank_dispatch_physical_map contains unassigned entries." + ) + + 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)) + choices = arr * quotient + r.sample(arr, k=remainder) + r.shuffle(choices) + return choices + + +@dataclass +class ModelConfigForExpertLocation: + num_layers: int + num_logical_experts: int + num_groups: int | None = None + + @staticmethod + def init_dummy(): + return ModelConfigForExpertLocation(num_layers=1, num_logical_experts=1) + + @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 ModelConfigForExpertLocation.init_dummy() + + +def compute_initial_expert_location_metadata( + server_args: ServerArgs, model_config: ModelConfig +) -> ExpertLocationMetadata: + data = server_args.init_expert_location + if data == "trivial": + return ExpertLocationMetadata.init_trivial(server_args, model_config) + + 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/tokenspeed/runtime/multimodal/__init__.py b/python/tokenspeed/runtime/multimodal/__init__.py new file mode 100644 index 0000000..87a1ace --- /dev/null +++ b/python/tokenspeed/runtime/multimodal/__init__.py @@ -0,0 +1,26 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Multimodal (VLM) runtime support. + +Per-item content hashing, pad-value substitution for MM-aware prefix caching, +SHM pixel transport between gateway and worker, the budget-bucketed encoder +CUDA-graph wrapper, and the M-RoPE helper. +""" diff --git a/python/tokenspeed/runtime/multimodal/embedder.py b/python/tokenspeed/runtime/multimodal/embedder.py new file mode 100644 index 0000000..4304f7a --- /dev/null +++ b/python/tokenspeed/runtime/multimodal/embedder.py @@ -0,0 +1,547 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Assemble LM input embeddings with multimodal encoder tokens spliced in. + +Three sequential phases: + + 1. ``_plan`` walks the active multimodal inputs in the current forward + batch and emits an :class:`EncodePlan` listing (a) the unique items + that still need to be encoded this iteration and (b) every flat + position in ``input_ids`` that should be filled from an encoder token, + along with the source range inside the owning item's encoded tensor. + + 2. ``_encode`` invokes the model-supplied encoder once per modality with + every miss in the batch in a single call, then writes each item's + output back onto the item itself (``item.encoded`` / + ``item.encoded_deepstack``). + + 3. ``_assemble`` runs the text-token embedding lookup and slices the + encoder-token ranges into the right positions using the plan's + :class:`ScatterRange` records. + +Per-item encoded tensors live on the :class:`MultimodalDataItem` itself, +not in an engine-global cache. Lifetime tracks the owning request: when +the request finishes and its ``RequestState`` is dropped, the tensors are +released by GC. Across chunked-prefill iterations of the same request the +item is identical Python object, so the second chunk sees ``item.encoded`` +already set and skips re-encoding. + +Within a single forward batch we still de-duplicate by modality and +``item.hash``: if two requests reference the same media content using +the same modality, only the first item is fed to the encoder; the second +request's scatter ranges read from the first item's ``encoded`` tensor. +""" + +from __future__ import annotations + +import logging +import time +from collections import defaultdict +from collections.abc import Callable +from dataclasses import dataclass, field +from typing import Any + +import torch +from torch import nn + +from tokenspeed.runtime.multimodal.inputs import ( + Modality, + MultimodalDataItem, + MultimodalForwardContext, + MultimodalInputs, +) +from tokenspeed.runtime.multimodal.shm_transport import ShmTensorHandle +from tokenspeed.runtime.utils.env import envs + +EncoderFn = Callable[[list[MultimodalDataItem]], torch.Tensor] + +logger = logging.getLogger(__name__) +LOG_MM_TIMING = envs.TOKENSPEED_LOG_MM_TIMING.get() + + +@dataclass +class EncoderSpec: + """Per-modality encoder registration. + + Bundles the encoder callable with whether its output needs to be + split into a main + deepstack pair via the model's + ``separate_deepstack_embeds`` hook. + """ + + fn: EncoderFn + deepstack: bool = False + + +# --------------------------------------------------------------------------- +# Input-id padding helper +# --------------------------------------------------------------------------- + + +def pad_input_tokens(input_ids: list[int], mm_inputs: MultimodalInputs) -> list[int]: + """Substitute placeholder token IDs with each item's ``pad_value``. + + The gateway produces ``input_ids`` with a single placeholder token + repeated across every multimodal-token position (e.g. ```` + repeated 1024 times for a 1024-token image). The prefix cache needs + each placeholder run to carry a content-derived ID so two different + images compare unequal. We rewrite each ``offsets`` range to the + item's pre-computed ``pad_value`` here. + """ + if not input_ids or not mm_inputs.mm_items: + return input_ids + + out = None + for item in mm_inputs.mm_items: + if item.pad_value is None or not item.offsets: + continue + if out is None: + out = list(input_ids) + pad_value = int(item.pad_value) + for offset_start, offset_end in item.offsets: + out[offset_start : offset_end + 1] = [pad_value] * ( + offset_end - offset_start + 1 + ) + return input_ids if out is None else out + + +# --------------------------------------------------------------------------- +# Plan structures +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class ScatterRange: + """One contiguous range to fill with multimodal encoder tokens. + + ``flat_dst_*`` are positions in the batch-flat ``input_ids`` tensor + (inclusive on both ends). ``item_src_*`` are positions within + ``item.encoded`` (also inclusive). ``item`` is the *canonical* item + holding the encoded tensor — for within-batch dedup'd entries it may + differ from the request-local item that produced the offsets. + """ + + flat_dst_start: int + flat_dst_end: int + item: MultimodalDataItem + item_src_start: int + item_src_end: int + + +@dataclass +class EncodePlan: + """Work to do this prefill iteration. + + ``misses_by_modality`` lists the canonical items the encoder needs to + process; each modality/content-hash pair appears at most once. + ``scatter_ranges`` describes every place an encoder token must land. + """ + + misses_by_modality: dict[Modality, list[MultimodalDataItem]] = field( + default_factory=lambda: defaultdict(list) + ) + scatter_ranges: list[ScatterRange] = field(default_factory=list) + aliases_by_canonical: dict[MultimodalDataItem, list[MultimodalDataItem]] = field( + default_factory=lambda: defaultdict(list) + ) + + def __bool__(self) -> bool: + return bool(self.scatter_ranges) + + +def _item_token_count(item: MultimodalDataItem) -> int: + """Total encoded tokens for an item. One offset per subgrid; the + encoder concatenates subgrid tokens in offsets order.""" + if not item.offsets: + return 0 + return sum(end - start + 1 for start, end in item.offsets) + + +# --------------------------------------------------------------------------- +# MultimodalEmbedder +# --------------------------------------------------------------------------- + + +class MultimodalEmbedder: + """Multimodal input embedding pipeline for one model executor.""" + + def __init__(self) -> None: + self._h2d_stream: torch.cuda.Stream | None = None + + # --- public entry point ------------------------------------------------ + + def apply( + self, + input_ids: torch.Tensor, + text_embedding: nn.Embedding, + ctx: MultimodalForwardContext | None, + encoders: dict[Modality, EncoderSpec], + multimodal_model: nn.Module, + is_decode_or_idle: bool = False, + ) -> tuple[torch.Tensor | None, dict[str, Any]]: + """Compose LM input embeddings with encoder tokens scattered in. + + Returns ``(None, {})`` when there is nothing multimodal to do this + forward (decode iteration, or no active multimodal inputs). The + caller falls back to the regular text-only path on that signal. + """ + if is_decode_or_idle or ctx is None or not ctx.has_extend_inputs(): + return None, {} + + total_started = time.perf_counter() if LOG_MM_TIMING else None + plan_started = time.perf_counter() if LOG_MM_TIMING else None + plan = self._plan(ctx) + plan_elapsed_ms = ( + (time.perf_counter() - plan_started) * 1000 + if plan_started is not None + else None + ) + if not plan: + return None, {} + + encode_started = time.perf_counter() if LOG_MM_TIMING else None + self._encode(plan, encoders, multimodal_model, input_ids.device) + encode_elapsed_ms = ( + (time.perf_counter() - encode_started) * 1000 + if encode_started is not None + else None + ) + + alias_started = time.perf_counter() if LOG_MM_TIMING else None + released_alias_features = self._share_encoded_aliases(plan) + alias_elapsed_ms = ( + (time.perf_counter() - alias_started) * 1000 + if alias_started is not None + else None + ) + + assemble_started = time.perf_counter() if LOG_MM_TIMING else None + input_embeds, kwargs = self._assemble( + input_ids, text_embedding, plan, encoders, multimodal_model + ) + assemble_elapsed_ms = ( + (time.perf_counter() - assemble_started) * 1000 + if assemble_started is not None + else None + ) + + cleanup_started = time.perf_counter() if LOG_MM_TIMING else None + released_encoded_features = self._drop_encoded_features(ctx) + cleanup_elapsed_ms = ( + (time.perf_counter() - cleanup_started) * 1000 + if cleanup_started is not None + else None + ) + if LOG_MM_TIMING and total_started is not None: + misses = { + modality.name: len(items) + for modality, items in plan.misses_by_modality.items() + if items + } + logger.info( + "mm_timing multimodal_embedder_apply_ms total=%.3f plan=%.3f " + "encode=%.3f alias=%.3f assemble=%.3f feature_cleanup=%.3f " + "scatter_ranges=%d misses=%s input_rows=%d aliases=%d " + "released_alias_features=%d released_encoded_features=%d", + (time.perf_counter() - total_started) * 1000, + plan_elapsed_ms, + encode_elapsed_ms, + alias_elapsed_ms, + assemble_elapsed_ms, + cleanup_elapsed_ms, + len(plan.scatter_ranges), + misses, + int(input_ids.numel()), + sum(len(items) for items in plan.aliases_by_canonical.values()), + released_alias_features, + released_encoded_features, + ) + return input_embeds, kwargs + + # --- phase 1: plan ----------------------------------------------------- + + def _plan(self, ctx: MultimodalForwardContext) -> EncodePlan: + plan = EncodePlan() + if not ctx.mm_inputs: + return plan + + # Within-batch dedup: first item per modality and content hash is + # canonical; duplicates reuse its encoded tensor. + canonical_by_key: dict[tuple[Modality, int], MultimodalDataItem] = {} + scheduled: set[MultimodalDataItem] = set() + + # Walk the FULL batch (including text-only / decode requests) + # so base offsets line up with the flat input_ids tensor that + # the caller hands us. Requests without mm input contribute + # nothing but still advance ``base``. + base = 0 + for req_idx, mm_inputs in enumerate(ctx.mm_inputs): + if req_idx >= len(ctx.extend_seq_lens) or req_idx >= len( + ctx.extend_prefix_lens + ): + break + seq = ctx.extend_seq_lens[req_idx] + if mm_inputs is None or seq <= 0: + base += max(seq, 0) + continue + + prefix = ctx.extend_prefix_lens[req_idx] + chunk_start = prefix + chunk_end_inc = prefix + seq - 1 + + for item in mm_inputs.mm_items: + if item is None or not item.offsets: + continue + + if item.encoded is not None: + canonical = item + elif ( + item.hash is not None + and (item.modality, item.hash) in canonical_by_key + ): + canonical = canonical_by_key[(item.modality, item.hash)] + else: + canonical = item + if item.hash is not None: + canonical_by_key[(item.modality, item.hash)] = item + + if canonical is not item: + plan.aliases_by_canonical[canonical].append(item) + + # src_cursor: start of current subgrid inside item.encoded. + src_cursor = 0 + for offset_start, offset_end in item.offsets: + span = offset_end - offset_start + 1 + overlap_start = max(offset_start, chunk_start) + overlap_end = min(offset_end, chunk_end_inc) + if overlap_start > overlap_end: + src_cursor += span + continue + + plan.scatter_ranges.append( + ScatterRange( + flat_dst_start=base + (overlap_start - prefix), + flat_dst_end=base + (overlap_end - prefix), + item=canonical, + item_src_start=src_cursor + (overlap_start - offset_start), + item_src_end=src_cursor + (overlap_end - offset_start), + ) + ) + if canonical.encoded is None and canonical not in scheduled: + scheduled.add(canonical) + plan.misses_by_modality[canonical.modality].append(canonical) + src_cursor += span + + base += seq + + return plan + + # --- phase 2: encode --------------------------------------------------- + + def _encode( + self, + plan: EncodePlan, + encoders: dict[Modality, EncoderSpec], + multimodal_model: nn.Module, + device: torch.device, + ) -> None: + for modality, items in plan.misses_by_modality.items(): + if not items: + continue + spec = encoders.get(modality) + if spec is None: + raise RuntimeError( + f"MultimodalEmbedder: no encoder registered for {modality}" + ) + + move_started = time.perf_counter() if LOG_MM_TIMING else None + self._move_features_to_device(items, device) + move_elapsed_ms = ( + (time.perf_counter() - move_started) * 1000 + if move_started is not None + else None + ) + encoder_started = time.perf_counter() if LOG_MM_TIMING else None + output = spec.fn(items) + if LOG_MM_TIMING and device.type == "cuda": + torch.cuda.synchronize(device) + encoder_elapsed_ms = ( + (time.perf_counter() - encoder_started) * 1000 + if encoder_started is not None + else None + ) + output = output.reshape(-1, output.shape[-1]) + + per_item_lens = [_item_token_count(it) for it in items] + per_item_embs = torch.split(output, per_item_lens, dim=0) + + if spec.deepstack: + for item, emb in zip(items, per_item_embs): + main, deep = multimodal_model.separate_deepstack_embeds(emb) + item.encoded = main + item.encoded_deepstack = deep + else: + for item, emb in zip(items, per_item_embs): + item.encoded = emb + if LOG_MM_TIMING: + logger.info( + "mm_timing encoder_ms modality=%s items=%d " + "encoder_output_tokens=%d move_h2d=%.3f encode=%.3f " + "per_item_tokens=%s", + modality.name, + len(items), + int(output.shape[0]), + move_elapsed_ms, + encoder_elapsed_ms, + per_item_lens, + ) + + def _share_encoded_aliases(self, plan: EncodePlan) -> int: + released = 0 + for canonical, aliases in plan.aliases_by_canonical.items(): + if canonical.encoded is None: + continue + for alias in aliases: + alias.encoded = canonical.encoded + alias.encoded_deepstack = canonical.encoded_deepstack + if self._drop_raw_feature(alias): + released += 1 + return released + + # --- phase 3: assemble ------------------------------------------------- + + def _assemble( + self, + input_ids: torch.Tensor, + text_embedding: nn.Embedding, + plan: EncodePlan, + encoders: dict[Modality, EncoderSpec], + multimodal_model: nn.Module, + ) -> tuple[torch.Tensor, dict[str, Any]]: + # Placeholder positions hold large content-derived IDs that exceed + # vocab_size; the lookup we run here is overwritten for those rows + # by the scatter below, but the lookup still needs valid indices. + vocab_size = text_embedding.num_embeddings + safe_ids = input_ids.clamp(min=0, max=vocab_size - 1) + input_embeds = text_embedding(safe_ids) + + kwargs: dict[str, Any] = {} + deepstack_buffer: torch.Tensor | None = None + deepstack_modalities = { + modality for modality, spec in encoders.items() if spec.deepstack + } + if any(r.item.modality in deepstack_modalities for r in plan.scatter_ranges): + num_deepstack = len(multimodal_model.deepstack_visual_indexes) + shape = input_embeds.shape[:-1] + (input_embeds.shape[-1] * num_deepstack,) + deepstack_buffer = torch.zeros( + shape, dtype=input_embeds.dtype, device=input_embeds.device + ) + kwargs["input_deepstack_embeds"] = deepstack_buffer + + for r in plan.scatter_ranges: + main = r.item.encoded + if main is None: + raise RuntimeError( + "MultimodalEmbedder: item scheduled for encode has no " + "encoded tensor after _encode; this is a bug" + ) + src = main[r.item_src_start : r.item_src_end + 1] + input_embeds[r.flat_dst_start : r.flat_dst_end + 1] = src.to( + dtype=input_embeds.dtype, device=input_embeds.device + ) + + if deepstack_buffer is not None and r.item.encoded_deepstack is not None: + deep_src = r.item.encoded_deepstack[ + r.item_src_start : r.item_src_end + 1 + ] + deepstack_buffer[r.flat_dst_start : r.flat_dst_end + 1] = deep_src.to( + dtype=input_embeds.dtype, device=input_embeds.device + ) + + return input_embeds, kwargs + + # --- device helpers ---------------------------------------------------- + + def _h2d_stream_on(self, device: torch.device) -> torch.cuda.Stream: + if self._h2d_stream is None: + self._h2d_stream = torch.cuda.Stream(device=device) + return self._h2d_stream + + def _move_features_to_device( + self, items: list[MultimodalDataItem], device: torch.device + ) -> None: + """Stage encoder features onto ``device`` on a dedicated H2D stream. + + Inputs that originate from the SHM transport are pinned, so the + H2D copy can actually run async with respect to the LM kernels + already queued on the current stream. We synchronise the current + stream with the H2D stream before returning so the encode call + sees the moved tensors. + """ + pending = [ + it + for it in items + if isinstance(it.feature, (torch.Tensor, ShmTensorHandle)) + and (isinstance(it.feature, ShmTensorHandle) or it.feature.device != device) + ] + if not pending: + return + + for it in pending: + if isinstance(it.feature, ShmTensorHandle): + it.feature = it.feature.consume() + + if device.type != "cuda": + for it in pending: + if isinstance(it.feature, torch.Tensor): + it.feature = it.feature.to(device, non_blocking=True) + return + + h2d = self._h2d_stream_on(device) + current = torch.cuda.current_stream(device) + with torch.cuda.stream(h2d): + for it in pending: + if isinstance(it.feature, torch.Tensor): + it.feature = it.feature.to(device, non_blocking=True) + current.wait_stream(h2d) + + @staticmethod + def _drop_raw_feature(item: MultimodalDataItem) -> bool: + if item.feature is None: + return False + if isinstance(item.feature, ShmTensorHandle): + item.feature.release() + item.feature = None + return True + + @staticmethod + def _drop_encoded_features(ctx: MultimodalForwardContext) -> int: + released = 0 + for mm in ctx.mm_inputs: + if mm is None: + continue + for it in mm.mm_items: + if it.encoded is not None and MultimodalEmbedder._drop_raw_feature(it): + released += 1 + return released + + +# Compatibility alias for model implementations that predate audio support. +VisionEmbedder = MultimodalEmbedder diff --git a/python/tokenspeed/runtime/multimodal/encoder_cudagraph.py b/python/tokenspeed/runtime/multimodal/encoder_cudagraph.py new file mode 100644 index 0000000..32449e0 --- /dev/null +++ b/python/tokenspeed/runtime/multimodal/encoder_cudagraph.py @@ -0,0 +1,643 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Budget-bucketed CUDA graph capture/replay for multimodal encoders. + +The :class:`EncoderCudaGraphWrapper` is the generic manager: it owns budget +selection, graph capture, replay buffer updates, greedy packing, and eager +fallback. Model/modality-specific input layout is supplied by an adapter object. +""" + +from __future__ import annotations + +import contextlib +from collections.abc import Callable +from dataclasses import dataclass +from functools import cached_property +from typing import Any, Protocol + +import torch + +from tokenspeed.runtime.distributed.comm_backend.registry import get_global_backend +from tokenspeed.runtime.utils import logger + + +@dataclass +class BudgetGraphMetadata: + """One captured budget graph. + + Replay copies a real batch into the captured input/metadata buffers, calls + ``graph.replay()``, then reads ``output_buffer``. + """ + + graph: torch.cuda.CUDAGraph + input_buffers: dict[str, torch.Tensor] + metadata_buffers: dict[str, torch.Tensor] + output_buffer: torch.Tensor + + +class EncoderCudaGraphBatch(Protocol): + """Minimal batch contract consumed by :class:`EncoderCudaGraphWrapper`.""" + + @property + def input_tensors(self) -> dict[str, torch.Tensor]: ... + + @property + def encoder_output_tokens(self) -> list[int]: ... + + @property + def metadata_sequences(self) -> list[int]: ... + + def num_items(self) -> int: ... + + def select(self, indices: list[int]) -> "EncoderCudaGraphBatch": ... + + +class EncoderCudaGraphAdapter(Protocol): + """Model/modality-specific contract used by :class:`EncoderCudaGraphWrapper`.""" + + @property + def modality_name(self) -> str: ... + + @property + def device(self) -> torch.device: ... + + @property + def dtype(self) -> torch.dtype: ... + + @property + def capture_tp_size(self) -> int: ... + + @property + def capture_tp_group(self) -> Any | None: ... + + def batch_from_items(self, items: list[Any]) -> EncoderCudaGraphBatch: ... + + def capture_batch_for_budget( + self, + encoder_output_token_budget: int, + max_batch_size: int, + metadata_sequence_budget: int, + device: torch.device, + dtype: torch.dtype, + ) -> EncoderCudaGraphBatch: ... + + def prepare_metadata( + self, + batch: EncoderCudaGraphBatch, + encoder_output_token_budget: int | None, + metadata_sequence_budget: int, + ) -> dict[str, Any]: ... + + def forward( + self, + input_tensors: dict[str, torch.Tensor], + metadata: dict[str, Any], + ) -> torch.Tensor: ... + + def postprocess( + self, + encoder_outs: list[torch.Tensor], + batch: EncoderCudaGraphBatch, + ) -> torch.Tensor: ... + + +@dataclass +class VisionEncoderBatch: + """Qwen/Kimi-style vision batch. + + Per-item patch rows are concatenated on dim 0 and indexed by ``grid_thw``. + ``.tolist()`` syncs are confined here, in the eager region, never inside a + graph replay. + """ + + tokens: torch.Tensor + grid: torch.Tensor + out_div: int + + @property + def input_tensors(self) -> dict[str, torch.Tensor]: + return {"tokens": self.tokens} + + @cached_property + def _grid_rows(self) -> list[list[int]]: + return self.grid.tolist() + + def num_items(self) -> int: + return self.grid.shape[0] + + @cached_property + def encoder_output_tokens(self) -> list[int]: + return [(t * h * w) // self.out_div for t, h, w in self._grid_rows] + + @cached_property + def cu_input(self) -> list[int]: + cu = [0] + for t, h, w in self._grid_rows: + cu.append(cu[-1] + t * h * w) + return cu + + @cached_property + def metadata_sequences(self) -> list[int]: + return [t for t, _, _ in self._grid_rows] + + def select(self, indices: list[int]) -> "VisionEncoderBatch": + """Sub-batch at ``indices``, preserving order.""" + cu = self.cu_input + if indices: + rows = torch.cat( + [ + torch.arange(cu[i], cu[i + 1], device=self.tokens.device) + for i in indices + ] + ) + else: + rows = torch.zeros(0, dtype=torch.long, device=self.tokens.device) + return VisionEncoderBatch(self.tokens[rows], self.grid[indices], self.out_div) + + +@dataclass +class VisionEncoderCudaGraphAdapter: + """Adapter for Qwen/Kimi-style ``grid_thw`` vision encoders.""" + + tower: Any + pre_encode: Callable[[list[Any]], tuple[torch.Tensor, torch.Tensor]] + post_encode: Callable[[list[torch.Tensor], torch.Tensor], torch.Tensor] + out_div: int + merge: int + input_feature_shape: tuple[int, ...] + modality_name: str = "vision" + out_squeeze_dim: int | None = None + capture_tp_size: int = 1 + capture_tp_group: Any | None = None + + @cached_property + def _param(self) -> torch.nn.Parameter: + return next(self.tower.parameters()) + + @property + def device(self) -> torch.device: + return self._param.device + + @property + def dtype(self) -> torch.dtype: + return self._param.dtype + + def synthetic_grid(self, encoder_output_token_budget: int) -> list[list[int]]: + n_patches = encoder_output_token_budget * self.out_div + units = max(n_patches // (self.merge * self.merge), 1) + a = 1 << (units.bit_length() // 2) + while a > 1 and units % a != 0: + a >>= 1 + b = units // a + return [[1, a * self.merge, b * self.merge]] + + def pad_cu_seqlens( + self, metadata: dict[str, Any], metadata_sequence_budget: int + ) -> None: + cu = metadata["cu_seqlens"] + target = metadata_sequence_budget + 1 + if cu.shape[0] > target: + raise RuntimeError( + f"{self.modality_name} encoder cudagraph needs {cu.shape[0] - 1} " + f"metadata sequences, but the configured limit is " + f"{metadata_sequence_budget}" + ) + pad = target - cu.shape[0] + if pad > 0: + metadata["cu_seqlens"] = torch.cat([cu, cu[-1:].expand(pad)]) + + def batch_from_items(self, items: list[Any]) -> VisionEncoderBatch: + tokens, grid = self.pre_encode(items) + return VisionEncoderBatch(tokens, grid, self.out_div) + + def capture_batch_for_budget( + self, + encoder_output_token_budget: int, + _max_batch_size: int, + _metadata_sequence_budget: int, + capture_device: torch.device, + capture_dtype: torch.dtype, + ) -> VisionEncoderBatch: + grid = torch.tensor( + self.synthetic_grid(encoder_output_token_budget), + device=capture_device, + dtype=torch.int32, + ) + tokens = torch.zeros( + (encoder_output_token_budget * self.out_div, *self.input_feature_shape), + device=capture_device, + dtype=capture_dtype, + ) + return VisionEncoderBatch(tokens, grid, self.out_div) + + def prepare_metadata( + self, + batch: EncoderCudaGraphBatch, + encoder_output_token_budget: int | None, + metadata_sequence_budget: int, + ) -> dict[str, Any]: + if not isinstance(batch, VisionEncoderBatch): + raise TypeError( + f"{self.modality_name} encoder cudagraph expected " + f"VisionEncoderBatch, got {type(batch).__name__}" + ) + metadata = dict(self.tower.prepare_metadata(batch.grid)) + if encoder_output_token_budget is not None: + self.pad_cu_seqlens(metadata, metadata_sequence_budget) + # Non-tensor scalar gets baked at capture. Use the per-budget worst + # case so replay never exceeds the captured attention max seqlen. + metadata["max_seqlen"] = encoder_output_token_budget * self.out_div + return metadata + + def forward( + self, input_tensors: dict[str, torch.Tensor], metadata: dict[str, Any] + ) -> torch.Tensor: + out = self.tower.forward_blocks(input_tensors["tokens"], metadata) + if self.out_squeeze_dim is not None: + out = out.squeeze(self.out_squeeze_dim) + return out + + def postprocess( + self, encoder_outs: list[torch.Tensor], batch: EncoderCudaGraphBatch + ) -> torch.Tensor: + if not isinstance(batch, VisionEncoderBatch): + raise TypeError( + f"{self.modality_name} encoder cudagraph expected " + f"VisionEncoderBatch, got {type(batch).__name__}" + ) + return self.post_encode(encoder_outs, batch.grid) + + +class EncoderCudaGraphWrapper: + """Generic budget-based CUDA graph manager for encoder callables. + + The wrapper does not know about image/video/audio internals. It only expects + batches to expose input tensors, per-item output row counts, per-item + metadata sequence counts, and a ``select`` operation. The adapter constructs + batches, prepares metadata, runs the captured forward, and postprocesses + output slices. + """ + + def __init__( + self, + *, + adapter: EncoderCudaGraphAdapter, + budget_range: tuple[int, int], + max_batch_size: int | None = None, + max_metadata_sequences_per_batch: int | None = None, + metadata_sequence_budget_from_encoder_output_budget: bool = False, + ): + self.adapter = adapter + self.device = adapter.device + self.dtype = adapter.dtype + self.modality_name = adapter.modality_name + + min_budget, max_budget = budget_range + self.encoder_output_token_budgets = self._generate_budgets( + min_budget, max_budget + ) + self.max_batch_size = ( + max_batch_size + if max_batch_size is not None + else max(1, max_budget // max(1, min_budget)) + ) + self.max_metadata_sequences_per_batch = max_metadata_sequences_per_batch + self.metadata_sequence_budget_from_encoder_output_budget = ( + metadata_sequence_budget_from_encoder_output_budget + ) + + self.capture_tp_size = adapter.capture_tp_size + self.capture_tp_group = adapter.capture_tp_group + + self.budget_graphs: dict[int, BudgetGraphMetadata] = {} + + metadata_sequence_budget_log_value = ( + self.max_metadata_sequences_per_batch + if self.max_metadata_sequences_per_batch is not None + else ( + "encoder_output_token_budget" + if self.metadata_sequence_budget_from_encoder_output_budget + else "batch" + ) + ) + logger.info( + "EncoderCudaGraphWrapper initialized: modality=%s, budgets=%s, " + "max_batch_size=%d, max_metadata_sequences_per_batch=%s, encoder_tp=%d", + self.modality_name, + self.encoder_output_token_budgets, + self.max_batch_size, + metadata_sequence_budget_log_value, + self.capture_tp_size, + ) + + def __call__(self, items: list[Any]) -> torch.Tensor: + batch = self.adapter.batch_from_items(items) + if not self.budget_graphs: + self.capture() + encoder_outs = self._dispatch(batch) + return self.adapter.postprocess(encoder_outs, batch) + + @staticmethod + def _generate_budgets(min_budget: int, max_budget: int) -> list[int]: + """Power-of-2 budgets in ``[min_budget, max_budget]``.""" + budgets: list[int] = [] + b = max(1, min_budget) + while b <= max_budget: + budgets.append(b) + b *= 2 + if not budgets or budgets[-1] < max_budget: + budgets.append(max_budget) + return budgets + + def _metadata_sequence_budget_for_encoder_output_budget( + self, encoder_output_token_budget: int + ) -> int: + if self.max_metadata_sequences_per_batch is not None: + return min( + self.max_metadata_sequences_per_batch, encoder_output_token_budget + ) + if self.metadata_sequence_budget_from_encoder_output_budget: + return encoder_output_token_budget + return self.max_batch_size + + def capture(self) -> None: + for encoder_output_token_budget in self.encoder_output_token_budgets: + self._capture_one(encoder_output_token_budget) + logger.info( + "Encoder CUDA graph capture complete: modality=%s, %d budget graphs.", + self.modality_name, + len(self.budget_graphs), + ) + + def _capture_one(self, encoder_output_token_budget: int) -> None: + metadata_sequence_budget = ( + self._metadata_sequence_budget_for_encoder_output_budget( + encoder_output_token_budget + ) + ) + batch = self.adapter.capture_batch_for_budget( + encoder_output_token_budget, + self.max_batch_size, + metadata_sequence_budget, + self.device, + self.dtype, + ) + metadata = dict( + self.adapter.prepare_metadata( + batch, encoder_output_token_budget, metadata_sequence_budget + ) + ) + input_buffers = batch.input_tensors + + # Warmup also forces lazy JIT / autotune before capture. + with torch.inference_mode(): + output = self.adapter.forward(input_buffers, metadata) + output_buffer = torch.empty_like(output) + + # Encoder TP > 1: capture must record per-layer all-reduce under the + # custom-AR capture context. + if self.capture_tp_size > 1 and self.capture_tp_group is not None: + ar_ctx: Any = get_global_backend().custom_ar.capture(self.capture_tp_group) + else: + ar_ctx = contextlib.nullcontext() + + # No pool= argument: each budget graph gets its own private pool. A + # shared pool collides custom-AR IPC registrations across budgets. + graph = torch.cuda.CUDAGraph() + with torch.inference_mode(), ar_ctx, torch.cuda.graph(graph): + output = self.adapter.forward(input_buffers, metadata) + output_buffer.copy_(output) + + # Only tensor entries are captured. Ints / None are baked at capture. + metadata_buffers = { + k: v for k, v in metadata.items() if isinstance(v, torch.Tensor) + } + self.budget_graphs[encoder_output_token_budget] = BudgetGraphMetadata( + graph=graph, + input_buffers=input_buffers, + metadata_buffers=metadata_buffers, + output_buffer=output_buffer, + ) + logger.debug( + "Captured encoder cudagraph: modality=%s, budget=%d, " + "max_batch_size=%d, metadata_sequence_budget=%d, buffers=%s", + self.modality_name, + encoder_output_token_budget, + self.max_batch_size, + metadata_sequence_budget, + {k: (v.dtype, tuple(v.shape)) for k, v in metadata_buffers.items()}, + ) + + def _smallest_fitting_budget( + self, total_encoder_output_tokens: int, total_metadata_sequences: int + ) -> int | None: + for budget in self.encoder_output_token_budgets: + if ( + budget >= total_encoder_output_tokens + and total_metadata_sequences + <= self._metadata_sequence_budget_for_encoder_output_budget(budget) + ): + return budget + return None + + @staticmethod + def _scatter_output_slices( + output: torch.Tensor, + indices: list[int], + per_item_encoder_output_tokens: list[int], + dest: dict[int, torch.Tensor], + clone: bool = False, + ) -> None: + """Slice ``output`` and scatter into ``dest`` by original item index.""" + offset = 0 + for idx in indices: + n_tokens = per_item_encoder_output_tokens[idx] + sliced = output[offset : offset + n_tokens] + dest[idx] = sliced.clone() if clone else sliced + offset += n_tokens + + def _run_budget_graph( + self, + batch: EncoderCudaGraphBatch, + encoder_output_token_budget: int, + ) -> torch.Tensor: + """Copy the batch into captured buffers, replay, and return output.""" + graph_meta = self.budget_graphs[encoder_output_token_budget] + metadata_sequence_budget = ( + self._metadata_sequence_budget_for_encoder_output_budget( + encoder_output_token_budget + ) + ) + + src_buffers = batch.input_tensors + if src_buffers.keys() != graph_meta.input_buffers.keys(): + raise RuntimeError( + f"{self.modality_name} encoder cudagraph input keys changed: " + f"capture={sorted(graph_meta.input_buffers.keys())}, " + f"replay={sorted(src_buffers.keys())}" + ) + + for key, buf in graph_meta.input_buffers.items(): + src = src_buffers[key] + n = src.shape[0] + if n > buf.shape[0]: + raise RuntimeError( + f"{self.modality_name} encoder cudagraph input {key} has " + f"{n} rows, but budget {encoder_output_token_budget} only " + f"captured {buf.shape[0]} rows" + ) + if src.shape[1:] != buf.shape[1:]: + raise RuntimeError( + f"{self.modality_name} encoder cudagraph input {key} " + f"shape changed after dim0: capture={tuple(buf.shape)}, " + f"replay={tuple(src.shape)}" + ) + buf.zero_() + buf[:n].copy_(src) + + metadata = dict( + self.adapter.prepare_metadata( + batch, encoder_output_token_budget, metadata_sequence_budget + ) + ) + replay_buffers = { + k: v for k, v in metadata.items() if isinstance(v, torch.Tensor) + } + + if replay_buffers.keys() != graph_meta.metadata_buffers.keys(): + raise RuntimeError( + f"{self.modality_name} encoder cudagraph metadata keys changed: " + f"capture={sorted(graph_meta.metadata_buffers.keys())}, " + f"replay={sorted(replay_buffers.keys())}" + ) + + for key, buf in graph_meta.metadata_buffers.items(): + new = replay_buffers[key] + if new.ndim == 0: + buf.copy_(new) + else: + if new.shape[1:] != buf.shape[1:]: + raise RuntimeError( + f"{self.modality_name} encoder cudagraph metadata {key} " + f"shape changed after dim0: capture={tuple(buf.shape)}, " + f"replay={tuple(new.shape)}" + ) + if new.shape[0] > buf.shape[0]: + raise RuntimeError( + f"{self.modality_name} encoder cudagraph metadata {key} " + f"has {new.shape[0]} rows, but the captured buffer only " + f"has {buf.shape[0]} rows" + ) + buf.zero_() + buf[: new.shape[0]].copy_(new) + + graph_meta.graph.replay() + return graph_meta.output_buffer + + def _run_eager(self, batch: EncoderCudaGraphBatch) -> torch.Tensor: + metadata = dict( + self.adapter.prepare_metadata( + batch, None, max(1, sum(batch.metadata_sequences)) + ) + ) + return self.adapter.forward(batch.input_tensors, metadata) + + def _dispatch(self, batch: EncoderCudaGraphBatch) -> list[torch.Tensor]: + """Greedy smallest-first pack into budget graphs with eager fallback.""" + num_items = batch.num_items() + max_budget = self.encoder_output_token_budgets[-1] + max_metadata_sequence_budget = ( + self._metadata_sequence_budget_for_encoder_output_budget(max_budget) + ) + per_item_encoder_output_tokens = batch.encoder_output_tokens + per_item_metadata_sequences = batch.metadata_sequences + + sorted_indices = sorted( + range(num_items), key=lambda i: per_item_encoder_output_tokens[i] + ) + + batches: list[tuple[list[int], int | None]] = [] + current_batch: list[int] = [] + current_batch_encoder_output_tokens = 0 + current_batch_metadata_sequences = 0 + for orig_idx in sorted_indices: + item_encoder_output_tokens = per_item_encoder_output_tokens[orig_idx] + item_metadata_sequences = per_item_metadata_sequences[orig_idx] + if ( + current_batch_encoder_output_tokens + item_encoder_output_tokens + <= max_budget + and len(current_batch) < self.max_batch_size + and current_batch_metadata_sequences + item_metadata_sequences + <= max_metadata_sequence_budget + ): + current_batch.append(orig_idx) + current_batch_encoder_output_tokens += item_encoder_output_tokens + current_batch_metadata_sequences += item_metadata_sequences + else: + if current_batch: + batches.append( + ( + current_batch, + self._smallest_fitting_budget( + current_batch_encoder_output_tokens, + current_batch_metadata_sequences, + ), + ) + ) + current_batch = [orig_idx] + current_batch_encoder_output_tokens = item_encoder_output_tokens + current_batch_metadata_sequences = item_metadata_sequences + if current_batch: + batches.append( + ( + current_batch, + self._smallest_fitting_budget( + current_batch_encoder_output_tokens, + current_batch_metadata_sequences, + ), + ) + ) + + # Packing reorders; restore original order before return. + outputs_by_orig_idx: dict[int, torch.Tensor] = {} + for batch_orig_indices, encoder_output_token_budget in batches: + sub_batch = batch.select(batch_orig_indices) + if encoder_output_token_budget is None: + with torch.inference_mode(): + raw = self._run_eager(sub_batch) + self._scatter_output_slices( + raw, + batch_orig_indices, + per_item_encoder_output_tokens, + outputs_by_orig_idx, + ) + else: + output = self._run_budget_graph(sub_batch, encoder_output_token_budget) + # clone: output is the shared, reused output_buffer. + self._scatter_output_slices( + output, + batch_orig_indices, + per_item_encoder_output_tokens, + outputs_by_orig_idx, + clone=True, + ) + + return [outputs_by_orig_idx[i] for i in range(num_items)] diff --git a/python/tokenspeed/runtime/multimodal/hash.py b/python/tokenspeed/runtime/multimodal/hash.py new file mode 100644 index 0000000..462eeb4 --- /dev/null +++ b/python/tokenspeed/runtime/multimodal/hash.py @@ -0,0 +1,84 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Content hashing for multimodal features. + +A multimodal feature -- an image/video pixel tensor, a numpy array, or a nested +list of them -- is folded into a single unsigned 64-bit integer. The runtime +uses that integer for within-batch dedup (duplicate features encode once) and +as the seed for a per-item pad value that substitutes the placeholder token ids +so the text-only prefix cache can prefix-match across requests. The hash only +needs to be deterministic and well distributed *within a run*: values are +computed once (in the gateway producer) and travel with the item, never +persisted or compared across builds, so the concrete digest is an +implementation detail. +""" + +import hashlib +from collections.abc import Iterable + +import numpy as np +import torch + +from tokenspeed.runtime.utils import flatten_nested_list + +# blake2b emits an 8-byte digest natively, which is exactly our key width. +_KEY_BYTES = 8 + +ByteChunk = bytes | bytearray | memoryview + + +def _fold(chunks: Iterable[ByteChunk]) -> int: + """Fold an ordered sequence of byte chunks into one unsigned 64-bit key.""" + digest = hashlib.blake2b(digest_size=_KEY_BYTES) + for chunk in chunks: + digest.update(chunk) + return int.from_bytes(digest.digest(), byteorder="big") + + +def _raw_bytes(buffer: torch.Tensor | np.ndarray) -> memoryview: + """Contiguous byte view of a tensor/array; CUDA tensors are pulled to host.""" + if isinstance(buffer, torch.Tensor): + if buffer.is_cuda: + buffer = buffer.cpu() + return memoryview(buffer.detach().contiguous().view(torch.uint8).numpy()) + return memoryview(np.ascontiguousarray(buffer)) + + +def hash_feature(feature) -> int: + """Deterministic unsigned 64-bit content hash of a multimodal feature. + + Handles a single tensor or numpy array, a (possibly nested) list of those, + and -- as a fallback -- any bytes-like or ``repr``-able object. + """ + if isinstance(feature, (torch.Tensor, np.ndarray)): + return _fold([_raw_bytes(feature)]) + + if isinstance(feature, list): + leaves = flatten_nested_list(feature) + if leaves and all(isinstance(x, (torch.Tensor, np.ndarray)) for x in leaves): + return _fold(_raw_bytes(x) for x in leaves) + # Non-array leaves (e.g. python scalars): hash a stable serialization. + return _fold([repr(tuple(leaves)).encode()]) + + if isinstance(feature, (bytes, bytearray, memoryview)): + return _fold([feature]) + + return _fold([repr(feature).encode()]) diff --git a/python/tokenspeed/runtime/multimodal/inputs.py b/python/tokenspeed/runtime/multimodal/inputs.py new file mode 100644 index 0000000..8fb611e --- /dev/null +++ b/python/tokenspeed/runtime/multimodal/inputs.py @@ -0,0 +1,191 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Multimodal request data structures used across processors and model adapters.""" + +from __future__ import annotations + +import dataclasses +import uuid +from enum import Enum, auto +from typing import Any + +import numpy as np +import torch + +from tokenspeed.runtime.multimodal.hash import hash_feature +from tokenspeed.runtime.multimodal.shm_transport import ShmTensorHandle +from tokenspeed.runtime.utils.env import envs + +# Multimodal pad-value substitute IDs: a placeholder mm token's id is rewritten +# to ``_MM_PAD_BASE + (hash & _MM_PAD_HASH_MASK)`` so duplicate features share +# the same substitute and prefix-match in the text-only prefix cache. The base +# sits well above any text vocab; the 30-bit mask keeps cross-hash collisions +# rare enough for long-running servers (~10^9 slots). +_MM_PAD_BASE = 1_000_000 +_MM_PAD_HASH_MASK = (1 << 30) - 1 + + +def is_mm_pad_value(token_ids: torch.Tensor) -> torch.Tensor: + """Bool mask of positions rewritten to a hash-derived multimodal pad id.""" + return (token_ids >= _MM_PAD_BASE) & (token_ids <= _MM_PAD_BASE + _MM_PAD_HASH_MASK) + + +def maybe_substitute_mm_pad( + input_ids: torch.Tensor, substitute_id: int | None +) -> torch.Tensor: + """Replace hash mm-pad positions with ``substitute_id``; no-op if None.""" + if substitute_id is None: + return input_ids + return input_ids.masked_fill(is_mm_pad_value(input_ids), substitute_id) + + +class Modality(Enum): + IMAGE = auto() + VIDEO = auto() + AUDIO = auto() + + +# ``eq=False`` on every dataclass below: tensor-valued fields crash the +# default element-wise ``__eq__`` and force ``__hash__`` to None. +@dataclasses.dataclass(eq=False) +class MultimodalDataItem: + modality: Modality + hash: int | None = None + pad_value: int | None = None + offsets: list | None = None + feature: torch.Tensor | np.ndarray | ShmTensorHandle | None = None + model_specific_data: dict[str, Any] = dataclasses.field(default_factory=dict) + # Encoder output for this item, populated on first encoder pass and reused + # across chunked-prefill iterations of the owning request. Lifetime is + # tied to the request: when the request finishes the item is GC'd and + # these tensors are released. ``encoded_deepstack`` is set only for + # deepstack-enabled modalities. + encoded: torch.Tensor | None = None + encoded_deepstack: torch.Tensor | None = None + # EPD (encode-prefill-decode): when set, this item's embedding is received + # from an encode worker over Mooncake into ``encoded`` instead of running the + # vision tower. A dict ``{bootstrap_room, bootstrap_host, bootstrap_port}`` + # naming the encode worker's rendezvous for this item's image (one room per + # item: the gateway splits the mm payload one item per image and the encode + # worker row-splits the concatenated-subgrid embedding per item). None for + # non-EPD items (left to the vision tower). + encode_handshake: dict | None = None + + def __getattr__(self, name: str): + if ( + "model_specific_data" in self.__dict__ + and name in self.__dict__["model_specific_data"] + ): + return self.__dict__["model_specific_data"][name] + raise AttributeError(f"'{self.__class__.__name__}' has no attribute '{name}'") + + def ensure_hash(self): + """Resolve ``self.hash`` to a concrete content id, lazily. + + The hash is resolved on demand rather than at construction because it + is usually supplied by the caller, a SHM-backed feature cannot be + hashed here without reading shared memory, and hashing inline bytes is + only worth doing once the value is actually needed. + + Resolution order: + * ``TOKENSPEED_MM_SKIP_COMPUTE_HASH`` -> a random id (dedup disabled); + * an already-set hash (e.g. the gateway-provided ``content_hash`` for + image/video) is kept as-is, no recompute; + * inline features the gateway does not hash (e.g. audio) are hashed + in-engine via ``hash_feature``; + * SHM-backed features must carry a caller-provided hash, else raise -- + we cannot hash a handle without reading shared memory. + """ + if envs.TOKENSPEED_MM_SKIP_COMPUTE_HASH.get(): + self.hash = uuid.uuid4().int + elif self.hash is None: + if isinstance(self.feature, ShmTensorHandle): + raise ValueError( + "SHM-backed multimodal items must carry content hash or " + "pad_value before TokenSpeed consumes them" + ) + self.hash = hash_feature(self.feature) + if self.hash is None: + raise RuntimeError("Failed to resolve multimodal item hash.") + + def set_pad_value(self): + if self.pad_value is not None: + return + self.ensure_hash() + self.pad_value = _MM_PAD_BASE + (self.hash & _MM_PAD_HASH_MASK) + + def is_modality(self, modality: Modality) -> bool: + return self.modality == modality + + +@dataclasses.dataclass(eq=False) +class MultimodalInputs: + mm_items: list[MultimodalDataItem] + im_token_id: int | None = None + video_token_id: int | None = None + mrope_positions: torch.Tensor | None = None + mrope_position_delta: torch.Tensor | None = None + mrope_position_delta_scalar: int | None = None + mrope_position_delta_repeated_cache: torch.Tensor | None = None + + def ensure_pad_values(self) -> None: + for item in self.mm_items: + item.set_pad_value() + + def publish_shm_features(self) -> None: + for item in self.mm_items: + if isinstance(item.feature, torch.Tensor): + item.feature = ShmTensorHandle.publish(item.feature) + + def attach_shm_features(self) -> None: + """Open every pending handle on this rank. Must run before the + cross-rank barrier in ``request_handler.recv_reqs``. + """ + for item in self.mm_items: + if isinstance(item.feature, ShmTensorHandle): + item.feature.attach() + + def release_shm_features(self) -> None: + for item in self.mm_items: + if isinstance(item.feature, ShmTensorHandle): + item.feature.release() + item.feature = None + + def has_pending_shm_features(self) -> bool: + return any(isinstance(item.feature, ShmTensorHandle) for item in self.mm_items) + + +@dataclasses.dataclass(eq=False) +class MultimodalForwardContext: + """Per-forward multimodal metadata for prefill embedding replacement.""" + + mm_inputs: list[MultimodalInputs | None] + extend_prefix_lens: list[int] + extend_seq_lens: list[int] + + def has_inputs(self) -> bool: + return bool(self.mm_inputs and any(x is not None for x in self.mm_inputs)) + + def has_extend_inputs(self) -> bool: + return any( + mm_input is not None and index < len(self.extend_seq_lens) + for index, mm_input in enumerate(self.mm_inputs) + ) diff --git a/python/tokenspeed/runtime/multimodal/mrope.py b/python/tokenspeed/runtime/multimodal/mrope.py new file mode 100644 index 0000000..449a28e --- /dev/null +++ b/python/tokenspeed/runtime/multimodal/mrope.py @@ -0,0 +1,323 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Multimodal RoPE (M-RoPE) position computation. + +The SMG gateway ships precomputed multimodal inputs but does not compute the +3-axis M-RoPE position_ids that MRoPE-aware models (the Qwen-VL family) need. +The engine computes them here on the un-padded input_ids, from the model config +plus the image/video ``grid_thw`` carried on the multimodal items. Non-MRoPE +models (e.g. Kimi-K2.5) return ``(None, None)``. + +This replaces the former per-model ``BaseMultimodalProcessor`` hierarchy + +``processor_registry``, whose only remaining live use after the SMG migration +was this single computation. +""" + +from __future__ import annotations + +import torch + +from tokenspeed.runtime.layers.rotary_embedding import MRotaryEmbedding + +# Architectures whose HF configs follow the Qwen-VL M-RoPE layout +# (vision_config.spatial_merge_size, image/video/vision_start token ids, etc.). +_MROPE_ARCHITECTURES = { + "Qwen3_5ForConditionalGeneration", + "Qwen3_5MoeForConditionalGeneration", + "Qwen3OmniMoeForConditionalGeneration", +} + +_QWEN3_OMNI_ARCHITECTURES = { + "Qwen3OmniMoeForConditionalGeneration", +} + + +def _modality_name(item) -> str: + modality = item.modality + return getattr(modality, "name", str(modality)).lower() + + +def _as_grid_rows(value, key: str) -> torch.Tensor: + grid = torch.as_tensor(value, dtype=torch.long).reshape(-1, 3).cpu() + if torch.any(grid <= 0): + raise ValueError(f"{key} must contain positive [T, H, W] values") + return grid + + +def _per_grid_seconds(item, num_grids: int, default: float) -> list[float]: + data = item.model_specific_data + value = data.get("video_second_per_grid") + if value is None: + return [default] * num_grids + + values = torch.as_tensor(value, dtype=torch.float32).flatten().tolist() + if len(values) == 1: + return values * num_grids + if len(values) != num_grids: + raise ValueError( + "video_second_per_grid must be scalar or have one value per video grid" + ) + return values + + +def _omni_media_segments(thinker_config, input_len: int, mm_items): + """Return authored media spans with one position descriptor per span.""" + merge = thinker_config.vision_config.spatial_merge_size + default_seconds = float(getattr(thinker_config, "seconds_per_chunk", 2.0)) + segments = [] + + for item in mm_items: + offsets = list(item.offsets or []) + if not offsets: + continue + modality = _modality_name(item) + + if modality == "audio": + descriptors = [("audio", None, None)] * len(offsets) + elif modality in ("image", "video"): + key = f"{modality}_grid_thw" + if key not in item.model_specific_data: + raise ValueError(f"Qwen3-Omni {modality} item is missing {key}") + if modality == "video" and "use_audio_in_video" in item.model_specific_data: + interleaved = torch.as_tensor( + item.model_specific_data["use_audio_in_video"] + ) + if bool(interleaved.any().item()): + raise ValueError( + "Qwen3-Omni use_audio_in_video=true is not supported" + ) + grids = _as_grid_rows(item.model_specific_data[key], key) + + if len(grids) != len(offsets): + raise ValueError( + f"Qwen3-Omni {modality} has {len(offsets)} placeholder spans " + f"but {len(grids)} grid rows" + ) + + seconds = ( + _per_grid_seconds(item, len(grids), default_seconds) + if modality == "video" + else [None] * len(grids) + ) + descriptors = [ + (modality, grid, second) + for grid, second in zip(grids, seconds, strict=True) + ] + else: + raise ValueError(f"Unsupported Qwen3-Omni modality: {modality}") + + for offset, descriptor in zip(offsets, descriptors, strict=True): + start, end = map(int, offset) + if start < 0 or end < start or end >= input_len: + raise ValueError( + f"Invalid Qwen3-Omni media offset [{start}, {end}] for " + f"input length {input_len}" + ) + segments.append((start, end, *descriptor)) + + segments.sort(key=lambda segment: segment[0]) + return segments, merge + + +def _compute_qwen3_omni_mrope_positions(hf_config, input_ids, mm_items): + """Compute Qwen3-Omni M-RoPE for independent image/audio/video inputs. + + The gateway's inclusive item offsets are authoritative. Audio advances all + three axes linearly. Vision uses T/H/W axes; video T is expressed on the + model's ``position_id_per_seconds`` clock. Audio extracted from a video is + intentionally not interleaved here (``use_audio_in_video=false``). + """ + thinker_config = getattr(hf_config, "thinker_config", hf_config) + position_id_per_seconds = float( + getattr(thinker_config, "position_id_per_seconds", 13) + ) + input_len = len(input_ids) + segments, spatial_merge_size = _omni_media_segments( + thinker_config, input_len, mm_items + ) + + position_chunks = [] + cursor = 0 + next_position = 0 + + def append_linear(length: int) -> None: + nonlocal next_position + if length <= 0: + return + positions = torch.arange( + next_position, next_position + length, dtype=torch.long + ) + position_chunks.append(positions.unsqueeze(0).expand(3, -1)) + next_position += length + + for start, end, modality, grid, seconds_per_grid in segments: + if start < cursor: + raise ValueError("Qwen3-Omni media placeholder spans overlap") + append_linear(start - cursor) + span = end - start + 1 + + if modality == "audio": + append_linear(span) + else: + t, h, w = (int(value) for value in grid) + if h % spatial_merge_size or w % spatial_merge_size: + raise ValueError( + "Qwen3-Omni vision grids must be divisible by spatial_merge_size" + ) + h //= spatial_merge_size + w //= spatial_merge_size + expected_span = t * h * w + if span != expected_span: + raise ValueError( + f"Qwen3-Omni {modality} placeholder span has {span} tokens; " + f"grid requires {expected_span}" + ) + + temporal = torch.arange(t, dtype=torch.float32) + if modality == "video": + temporal *= float(seconds_per_grid) * position_id_per_seconds + else: + temporal *= position_id_per_seconds + temporal = temporal.to(torch.long) + temporal = temporal.view(-1, 1).expand(-1, h * w).flatten() + height = ( + torch.arange(h, dtype=torch.long) + .view(1, -1, 1) + .expand(t, -1, w) + .flatten() + ) + width = ( + torch.arange(w, dtype=torch.long) + .view(1, 1, -1) + .expand(t, h, -1) + .flatten() + ) + media_positions = ( + torch.stack((temporal, height, width), dim=0) + next_position + ) + position_chunks.append(media_positions) + next_position = int(media_positions.max().item()) + 1 + + cursor = end + 1 + + append_linear(input_len - cursor) + positions = ( + torch.cat(position_chunks, dim=1) + if position_chunks + else torch.empty((3, 0), dtype=torch.long) + ) + if positions.shape[1] != input_len: + raise RuntimeError( + "Qwen3-Omni M-RoPE position count does not match the input length" + ) + delta = torch.tensor([[next_position - input_len]], dtype=torch.long) + return positions, delta + + +def compute_mrope_positions(hf_config, input_ids, mm_items): + """Compute ``(mrope_positions, mrope_position_delta)`` for MRoPE models. + + ``mm_items`` are the precomputed ``MultimodalDataItem``s (their + ``model_specific_data`` carries ``image_grid_thw`` / ``video_grid_thw``). + Returns ``(None, None)`` for non-MRoPE models. + """ + architectures = getattr(hf_config, "architectures", None) or [] + if not any(arch in _MROPE_ARCHITECTURES for arch in architectures): + return None, None + + if any(arch in _QWEN3_OMNI_ARCHITECTURES for arch in architectures): + return _compute_qwen3_omni_mrope_positions(hf_config, input_ids, mm_items) + + image_grids = [ + item.model_specific_data["image_grid_thw"] + for item in mm_items + if "image_grid_thw" in item.model_specific_data + ] + video_grids = [ + item.model_specific_data["video_grid_thw"] + for item in mm_items + if "video_grid_thw" in item.model_specific_data + ] + image_grid_thw = torch.cat(image_grids, dim=0) if image_grids else None + video_grid_thw = torch.cat(video_grids, dim=0) if video_grids else None + + # Qwen3.5 models compute M-RoPE with one video segment per temporal grid. + # The vision encoder still consumes the original grid [T, H, W], but the + # text prompt contains T separate <|video_pad|> runs. Split only the RoPE + # grid to match HuggingFace's Qwen3.5 get_rope_index behavior. + if video_grid_thw is not None and getattr(hf_config, "model_type", None) in ( + "qwen3_5", + "qwen3_5_moe", + ): + video_grid_thw = torch.repeat_interleave( + video_grid_thw, video_grid_thw[:, 0].to(torch.long), dim=0 + ).clone() + video_grid_thw[:, 0] = 1 + + input_ids_tensor = torch.tensor(input_ids, dtype=torch.long).unsqueeze(0) + mrope_positions, mrope_position_delta = MRotaryEmbedding.get_rope_index( + spatial_merge_size=hf_config.vision_config.spatial_merge_size, + image_token_id=hf_config.image_token_id, + video_token_id=hf_config.video_token_id, + vision_start_token_id=hf_config.vision_start_token_id, + model_type=hf_config.model_type, + tokens_per_second=getattr(hf_config.vision_config, "tokens_per_second", None), + input_ids=input_ids_tensor, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + ) + return mrope_positions.squeeze(1), mrope_position_delta + + +def extend_mrope_positions_for_retracted_request( + mrope_positions: torch.Tensor, output_ids_len: int +) -> torch.Tensor: + """Extend ``mrope_positions`` to cover already-generated output tokens. + + When a request carrying M-RoPE positions is retracted, the positions must be + extended over the output_ids generated so far. Output tokens are pure text, + so all three axes share the same incremental sequence. + + Args: + mrope_positions: original positions, shape ``(3, origin_input_ids_len)``. + output_ids_len: number of output tokens to generate positions for. + + Returns: + Extended positions, shape ``(3, origin_input_ids_len + output_ids_len)``. + """ + if output_ids_len <= 0: + return mrope_positions + + # Continue the incremental sequence from the last input position. + last_position = mrope_positions[:, -1] # (3,) + start_pos = last_position[0] + 1 + output_positions = ( + torch.arange( + start_pos, + start_pos + output_ids_len, + dtype=torch.int64, + device=mrope_positions.device, + ) + .unsqueeze(0) + .expand(3, -1) + ) # (3, output_ids_len) + + return torch.cat([mrope_positions, output_positions], dim=1) diff --git a/python/tokenspeed/runtime/multimodal/shm_transport.py b/python/tokenspeed/runtime/multimodal/shm_transport.py new file mode 100644 index 0000000..ed6ad24 --- /dev/null +++ b/python/tokenspeed/runtime/multimodal/shm_transport.py @@ -0,0 +1,164 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""POSIX SHM handle for cross-process multimodal feature tensors. + +The lifecycle keeps the unlink race-free for tensor-parallel ranks while still +allowing the model-side multimodal planner to deduplicate requests before any +large payload copy happens: + +``publish`` (producer) -> ``attach`` (every rank, before barrier) -> +``consume`` (only encoder misses) or ``release`` (deduplicated aliases). +""" + +from __future__ import annotations + +import logging +import time +from dataclasses import dataclass, field +from multiprocessing import shared_memory + +import torch + +from tokenspeed.runtime.utils.env import envs + +logger = logging.getLogger(__name__) +LOG_MM_TIMING = envs.TOKENSPEED_LOG_MM_TIMING.get() + + +@dataclass +class ShmTensorHandle: + """Pickle-safe handle to a CPU tensor in a POSIX SHM segment.""" + + shm_name: str + shape: tuple[int, ...] + dtype: torch.dtype + _segment: shared_memory.SharedMemory | None = field( + default=None, init=False, repr=False, compare=False + ) + + @classmethod + def publish(cls, tensor: torch.Tensor) -> ShmTensorHandle: + nbytes = tensor.numel() * tensor.element_size() + shm = shared_memory.SharedMemory(create=True, size=nbytes) + try: + shm_bytes = torch.frombuffer(shm.buf, dtype=torch.uint8) + shm_bytes.copy_(tensor.contiguous().view(torch.uint8).reshape(-1)) + except BaseException: + shm.close() + shm.unlink() + raise + name = shm.name + shm.close() + return cls(shm_name=name, shape=tuple(tensor.shape), dtype=tensor.dtype) + + def attach(self) -> None: + """Open the SHM segment on this rank. Must run before the cross-rank + barrier so unlink in ``consume()`` cannot race another rank's open. + """ + if self._segment is None: + self._segment = shared_memory.SharedMemory(name=self.shm_name) + + def consume(self) -> torch.Tensor: + """Copy into a pinned tensor (so downstream non_blocking H2D is real), + close this rank's FD, and unlink. ``attach()`` must have run. + """ + if self._segment is None: + raise RuntimeError( + f"ShmTensorHandle({self.shm_name!r}) must be attach()'d " + "before consume() (or has already been consumed on this rank)" + ) + segment = self._segment + started = time.perf_counter() if LOG_MM_TIMING else None + try: + dst = torch.empty(self.shape, dtype=self.dtype, pin_memory=True) + src = torch.frombuffer(segment.buf, dtype=self.dtype).reshape(self.shape) + dst.copy_(src) + finally: + self._segment = None + segment.close() + try: + segment.unlink() + except FileNotFoundError: + # Another rank already won the unlink race; benign. + pass + if LOG_MM_TIMING and started is not None: + logger.info( + "mm_timing shm_consume_ms name=%s elapsed=%.3f shape=%s dtype=%s", + self.shm_name, + (time.perf_counter() - started) * 1000, + list(self.shape), + self.dtype, + ) + return dst + + def release(self) -> None: + """Close and unlink a SHM segment without materializing the tensor.""" + started = time.perf_counter() if LOG_MM_TIMING else None + segment = self._segment + self._segment = None + try: + if segment is None: + segment = shared_memory.SharedMemory(name=self.shm_name) + segment.close() + try: + segment.unlink() + except FileNotFoundError: + pass + except FileNotFoundError: + pass + if LOG_MM_TIMING and started is not None: + logger.info( + "mm_timing shm_release_ms name=%s elapsed=%.3f shape=%s dtype=%s", + self.shm_name, + (time.perf_counter() - started) * 1000, + list(self.shape), + self.dtype, + ) + + +def sync_shm_features(reqs, group, group_size: int) -> None: + """Attach SHM-backed features in ``reqs`` on every rank. + + The barrier makes later consume/release unlink race-free in multi-rank + setups. Actual materialization is intentionally deferred until the + multimodal encoder planner has deduplicated the batch. + """ + pending = [ + mm + for req in reqs + if (mm := getattr(req, "multimodal_inputs", None)) is not None + and mm.has_pending_shm_features() + ] + if not pending: + return + started = time.perf_counter() if LOG_MM_TIMING else None + for mm in pending: + mm.attach_shm_features() + if group_size > 1: + torch.distributed.barrier(group) + if LOG_MM_TIMING and started is not None: + item_count = sum(len(mm.mm_items) for mm in pending) + logger.info( + "mm_timing shm_attach_ms requests=%d items=%d elapsed=%.3f", + len(pending), + item_count, + (time.perf_counter() - started) * 1000, + ) diff --git a/python/tokenspeed/runtime/pd/base/bootstrap.py b/python/tokenspeed/runtime/pd/base/bootstrap.py new file mode 100644 index 0000000..ad39bd2 --- /dev/null +++ b/python/tokenspeed/runtime/pd/base/bootstrap.py @@ -0,0 +1,195 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import asyncio +import logging +import threading +from dataclasses import dataclass + +from aiohttp import web + +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +@dataclass +class BootstrapInfo: + bootstrap_host: str + bootstrap_port: int + bootstrap_room: int + + +class DisaggBootstrapServerBase: + """HTTP rendezvous server shared by both PD and EPD transfer roles. + + A data-source rank PUTs its (ip, port) keyed by (dp_group, tp_rank_in_dp); + the peer GETs that back to open a Mooncake session, and uses the sentinel + ``engine_rank == -1`` query to sync the source's parallel sizes. The routing, + dp-group sharding, and server lifecycle are role-neutral and live here; + role-specific parallel-info fields (e.g. the KV path's MLA / kv-page lengths) + are layered in by subclasses via :meth:`_ingest_put_extra` (record extra PUT + fields) and :meth:`_extra_parallel_info` (add them to the GET sync response). + """ + + def __init__(self, port: int): + self.port = port + self.app = web.Application() + self.lock = asyncio.Lock() + self._setup_routes() + self.world_size = None + self.dp_size = None + self.tp_size_per_dp_rank = None + self.prefill_port_table: dict[int, dict[int, dict[str, str | int]]] = {} + + 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"] + world_size = data["world_size"] + dp_size = data["dp_size"] + rank_ip = data["rank_ip"] + rank_port = int(data["rank_port"]) + engine_rank = int(data["engine_rank"]) + self._ingest_put_extra(data) + + if self.world_size is None: + self.world_size = world_size + + if self.dp_size is None: + self.dp_size = dp_size + + tp_size_per_dp_rank = world_size // dp_size + if self.tp_size_per_dp_rank is None: + self.tp_size_per_dp_rank = tp_size_per_dp_rank + + if role == "Prefill": + dp_group = engine_rank // tp_size_per_dp_rank + tp_rank_in_dp_group = engine_rank % tp_size_per_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] = {} + + self.prefill_port_table[dp_group][tp_rank_in_dp_group] = { + "rank_ip": rank_ip, + "rank_port": rank_port, + } + logger.debug( + "Register prefill bootstrap: %s with rank_ip: %s and rank_port: %s", + engine_rank, + rank_ip, + 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.world_size, + "prefill_dp_size": self.dp_size, + } + prefill_parallel_info.update(self._extra_parallel_info()) + 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 _ingest_put_extra(self, data: dict) -> None: + """Record role-specific fields off a register PUT. Default: none.""" + + def _extra_parallel_info(self) -> dict: + """Role-specific fields to merge into the GET parallel-info sync response. + Default: none.""" + return {} + + 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, port=self.port) + self._loop.run_until_complete(site.start()) + self._loop.run_forever() + except Exception as exc: + logger.error("Server error: %s", str(exc)) + 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") diff --git a/python/tokenspeed/runtime/pd/base/manager.py b/python/tokenspeed/runtime/pd/base/manager.py new file mode 100644 index 0000000..6298a56 --- /dev/null +++ b/python/tokenspeed/runtime/pd/base/manager.py @@ -0,0 +1,101 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import threading +from functools import cache + +import zmq + +from tokenspeed.runtime.pd.base.status import TransferPoll + + +class DisaggManagerBase: + """Transfer-engine handle + ZMQ control socket + room-keyed status FSM shared + by the KV (prefill->decode) and embedding (encode->prefill) managers. + + Transport-neutral: the data-plane ``engine`` is constructed by the subclass + and injected here, so this base carries no vendor (Mooncake) binding -- only + the control-plane ZMQ socket and the status FSM. It carries no + KV/MLA/embedding-specific state either; subclasses add their role args/buffers + and implement :meth:`register_buffer_to_engine`. The subclass MUST set the + attributes that :meth:`register_buffer_to_engine` reads (e.g. ``kv_args`` / + ``embedding_args``) BEFORE calling ``super().__init__`` here, because this + constructor registers the buffers as its last engine step. + + ``request_status`` is the room-keyed FSM the senders/receivers poll; + :meth:`update_status` advances it monotonically and makes ``Failed`` sticky + so a late success cannot resurrect a broken transfer. + """ + + def __init__(self, *, engine): + self.engine = engine + self.server_socket = zmq.Context().socket(zmq.PULL) + self.register_buffer_to_engine() + + self.rank_port = None + self.request_status: dict[int, int] = {} + self.failure_records: dict[int, str] = {} + self.failure_lock = threading.Lock() + + def register_buffer_to_engine(self): + """Register this role's RDMA buffers with the engine. Abstract.""" + raise NotImplementedError + + def get_session_id(self) -> str: + return self.engine.get_session_id() + + @cache + def _connect(self, endpoint: str): + socket = zmq.Context().socket(zmq.PUSH) + socket.connect(endpoint) + return socket + + def check_status(self, bootstrap_room: int): + """Status of ``bootstrap_room``; raises ``KeyError`` if never seen.""" + return self.request_status[bootstrap_room] + + def room_status(self, bootstrap_room: int): + """Status of ``bootstrap_room``, or ``None`` if unknown -- the non-raising + read for lease/reap probes (an unknown or already-reaped room is a normal + answer rather than a ``KeyError``).""" + return self.request_status.get(bootstrap_room) + + def update_status(self, bootstrap_room: int, status: int): + if bootstrap_room not in self.request_status: + self.request_status[bootstrap_room] = status + else: + # Status is only allowed to be incremented unless either side has + # observed a failure. Failed is sticky so a late success cannot + # resurrect a broken transfer. + if ( + self.request_status[bootstrap_room] == TransferPoll.Failed + or status == TransferPoll.Failed + ): + self.request_status[bootstrap_room] = TransferPoll.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 diff --git a/python/tokenspeed/runtime/pd/base/mooncake_engine.py b/python/tokenspeed/runtime/pd/base/mooncake_engine.py new file mode 100644 index 0000000..b15043d --- /dev/null +++ b/python/tokenspeed/runtime/pd/base/mooncake_engine.py @@ -0,0 +1,154 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import logging + +logger = logging.getLogger(__name__) + + +class MooncakeTransferEngine: + def __init__(self, hostname: str, gpu_id: int, ib_device: str | None = 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 TokenSpeed 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"{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 initialize( + self, + hostname: str, + device_name: str | None, + ) -> None: + """Initialize the mooncake instance.""" + 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 transfer_submit_write( + self, session_id: str, buffer: int, peer_buffer_address: int, length: int + ) -> int: + """ASynchronously transfer data to the specified address.""" + + batch_id = self.engine.transfer_submit_write( + session_id, buffer, peer_buffer_address, length + ) + return batch_id + + def transfer_check_status(self, batch_id: int) -> int: + status = self.engine.transfer_check_status(batch_id) + return status + + def get_session_id(self): + return self.session_id diff --git a/python/tokenspeed/runtime/pd/base/status.py b/python/tokenspeed/runtime/pd/base/status.py new file mode 100644 index 0000000..ed0df07 --- /dev/null +++ b/python/tokenspeed/runtime/pd/base/status.py @@ -0,0 +1,28 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + + +class TransferPoll: + Failed = 0 + Bootstrapping = 1 + Bootstrapped = 2 + WaitingForInput = 3 + Transferring = 4 + Success = 5 diff --git a/python/tokenspeed/runtime/pd/decode_executor.py b/python/tokenspeed/runtime/pd/decode_executor.py new file mode 100644 index 0000000..cdc198a --- /dev/null +++ b/python/tokenspeed/runtime/pd/decode_executor.py @@ -0,0 +1,256 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + + +import numpy as np +import torch +from tokenspeed_scheduler import PD, Forward + +from tokenspeed.runtime.pd.base.bootstrap import BootstrapInfo +from tokenspeed.runtime.pd.base.status import TransferPoll +from tokenspeed.runtime.pd.mooncake.decode import MooncakeKVManagerDecode +from tokenspeed.runtime.pd.mooncake.receiver import MooncakeKVReceiver +from tokenspeed.runtime.pd.utils import ( + TransferBackend, + poll_and_all_reduce, +) +from tokenspeed.runtime.utils import get_colorful_logger +from tokenspeed.runtime.utils.dispatch import TypeBasedDispatcher + +logger = get_colorful_logger(__name__) + + +class DisaggDecodeExecutor: + def __init__( + self, backend: TransferBackend, args, kv_args, gloo_group, page_size: int + ): + self.transfer_backend = backend + self.bootstrap_port = args.bootstrap_port + self.page_size = page_size + self._dispatcher = TypeBasedDispatcher( + [ + (Forward.FlatForwardOp, self._prefill), + ] + ) + self.receivers: dict[int, MooncakeKVReceiver] = {} + self.kv_manager = MooncakeKVManagerDecode(args, kv_args) + self.gloo_group = gloo_group + self._local_states = {} + self._request_pool_indices: dict[str, int] = {} + self._remote_spec_candidate_ids: dict[str, tuple[int, list[int]]] = {} + + def _bootstrap(self, request_id, info): + self.receivers[request_id] = MooncakeKVReceiver( + mgr=self.kv_manager, + bootstrap_addr=f"{info.bootstrap_host}:{info.bootstrap_port}", + bootstrap_room=info.bootstrap_room, + ) + + @staticmethod + def _mamba_indices(op, index: int): + indices = getattr(op, "mamba_pool_indices", None) + if indices is None or index >= len(indices): + return None + slot = int(indices[index]) + if slot < 0: + return None + return np.array([slot], dtype=np.int64) + + @staticmethod + def _mamba_checkpoint_indices(op, index: int): + indices = getattr(op, "mamba_checkpoint_dst_indices", None) + if indices is None or index >= len(indices): + return None + slot = int(indices[index]) + if slot < 0: + return None + return np.array([slot], dtype=np.int64) + + @classmethod + def _mamba_transfer_indices(cls, op, index: int): + working = cls._mamba_indices(op, index) + if working is None: + return None + checkpoint = cls._mamba_checkpoint_indices(op, index) + if checkpoint is None: + return working + + slots = [int(x) for x in working.tolist()] + for slot in checkpoint.tolist(): + slot = int(slot) + if slot >= 0 and slot not in slots: + slots.append(slot) + return np.array(slots, dtype=np.int64) + + def _prefill(self, op): + logger.debug( + "[decode][_prefill] op: request_ids=%s occupied_pages=%s " + "begins=%s sizes=%s request_pool_indices=%s extend_prefix_lens=%s", + list(op.request_ids), + [list(p) for p in op.occupied_pages], + list(op.begins), + list(op.sizes), + list(op.request_pool_indices), + list(op.extend_prefix_lens), + ) + + for i, request_id in enumerate(op.request_ids): + if request_id not in self.receivers: + # Request failed and its receiver was cleaned up in generate_events; + # the scheduler may still dispatch its forward op one last time. + continue + extend_prefix_len = op.extend_prefix_lens[i] + kv_indices = np.array( + op.occupied_pages[i][extend_prefix_len // self.page_size :], + dtype=np.int64, + ) + aux_index = op.request_pool_indices[i] + mamba_indices = self._mamba_transfer_indices(op, i) + self._request_pool_indices[request_id] = aux_index + self.receivers[request_id].prefill( + kv_indices, + aux_index, + extend_prefix_len, + None, # mla_l1_5_args + mamba_indices, + ) + + def register(self, request_id: str, bootstrap_info: BootstrapInfo): + self._local_states[request_id] = TransferPoll.Bootstrapping + self._bootstrap(request_id, bootstrap_info) + + def execute(self, op): + if not isinstance(op, Forward.FlatForwardOp): + raise TypeError(f"Expected FlatForwardOp, got {type(op).__name__}.") + self._dispatcher(op) + + def generate_events(self): + if not self.receivers: + return [] + polls = poll_and_all_reduce(self.receivers.values(), self.gloo_group) + + events = [] + to_remove = [] + for req_id, poll in zip(list(self.receivers.keys()), polls): + if ( + self._local_states[req_id] == TransferPoll.Bootstrapping + and poll == TransferPoll.Bootstrapped + ): + logger.debug( + "[decode][generate_events] rid=%s -> BootstrappedEvent", req_id + ) + events.append(PD.BootstrappedEvent(req_id)) + self._local_states[req_id] = TransferPoll.Bootstrapped + elif poll == TransferPoll.Failed: + logger.warning( + "[decode][generate_events] rid=%s -> FailedEvent", req_id + ) + events.append(PD.FailedEvent(req_id)) + # Drop the failed receiver so it is not polled again. Without this + # a single failed request keeps re-emitting FailedEvent every loop + # (poll stays Failed), wedging the whole conn-1 scheduler. + to_remove.append(req_id) + elif ( + self._local_states[req_id] == TransferPoll.Bootstrapped + and poll == TransferPoll.Success + ): + # Read bootstrap_token from the ZMQ-delivered table in kv_manager. + # The decode_thread stored it there when it received the Success status + # message from the prefill side. bootstrap_room == bootstrap_info.bootstrap_room, + # which is the key used in MooncakeKVReceiver. + self._local_states[req_id] = TransferPoll.Success + bootstrap_room = self.receivers[req_id].bootstrap_room + bootstrap_token, spec_candidate_ids = ( + self.kv_manager.pop_prefill_metadata(bootstrap_room) + ) + receiver = self.receivers[req_id] + if ( + spec_candidate_ids is not None + and req_id in self._request_pool_indices + and getattr( + receiver, + "supports_remote_spec_candidates", + True, + ) + ): + self._remote_spec_candidate_ids[req_id] = ( + self._request_pool_indices[req_id], + spec_candidate_ids, + ) + logger.debug( + "[decode][generate_events] rid=%s -> RemotePrefillDoneEvent bootstrap_token=%s", + req_id, + bootstrap_token, + ) + # Use RemotePrefillDoneEvent to carry the bootstrap_token to event_loop; + # the C++ FSM will extend it into the TokenContainer via + # fsm::RemotePrefillDoneEvent::operator()(Prefilling&&). + event = PD.RemotePrefillDoneEvent( + req_id, bootstrap_token if bootstrap_token != -1 else -1 + ) + events.append(event) + to_remove.append(req_id) + else: + pass + for req_id in to_remove: + # Best-effort cleanup mirroring prefill side; request_id is stable + # so without explicit pop these dicts would grow unbounded across + # failed requests. NOTE: _remote_spec_candidate_ids must NOT be + # popped here — its consumer pop_remote_spec_candidate_ids runs + # later inside event_loop._process_kv_transfer_events, after we return. + # That dict is small (one tuple per Success request, between + # generate_events emitting RemotePrefillDoneEvent and event_loop + # consuming it) and is naturally drained by the pop path; an + # eager pop here drops the spec candidates on the floor and the + # next decode forward reads uninitialized future_input_map tail, + # causing CUDA illegal memory access on embedding lookup. + self.receivers.pop(req_id, None) + self._request_pool_indices.pop(req_id, None) + self._local_states.pop(req_id, None) + + return events + + def pop_remote_spec_candidate_ids(self, request_id: str): + return self._remote_spec_candidate_ids.pop(request_id, None) + + def reset_valid_cache_length( + self, forward_op, runtime_states, execution_stream, device + ): + num_extends = forward_op.num_extends() + extend_request_pool_indices = torch.tensor( + forward_op.request_pool_indices[:num_extends], + dtype=torch.int64, + device="cpu", + pin_memory=True, + ).to(device, non_blocking=True) + extend_prefix_lens = torch.tensor( + forward_op.prefill_lengths[:num_extends], + dtype=torch.int32, + device="cpu", + pin_memory=True, + ).to(device, non_blocking=True) + # HostTodevice segment ends + + execution_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(execution_stream): + if num_extends > 0: + runtime_states.reset_states( + extend_request_pool_indices, extend_prefix_lens + ) diff --git a/python/tokenspeed/runtime/pd/disaggregation_decode_scheduler.py b/python/tokenspeed/runtime/pd/disaggregation_decode_scheduler.py new file mode 100755 index 0000000..48aa9cb --- /dev/null +++ b/python/tokenspeed/runtime/pd/disaggregation_decode_scheduler.py @@ -0,0 +1,179 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the FluentLLM project +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import logging +from typing import TYPE_CHECKING + +import torch + +from tokenspeed.runtime.cache.req_to_token_pool import ReqToTokenPoolInfo +from tokenspeed.runtime.execution.forward_batch_info import ForwardMode +from tokenspeed.runtime.sampling.sampling_batch_info import SamplingBatchInfo + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from tokenspeed.runtime.configs.model_config import ModelConfig + from tokenspeed.runtime.engine.schedule_batch import ScheduleBatch + from tokenspeed.runtime.utils.server_args import ServerArgs + + +class DisaggDecodeScheduler: + + def prepare_for_prebuilt_extend(self: ScheduleBatch): + """ + Prepare a prebuilt extend by populate metadata + """ + + 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 + ] + if offset + req.extend_input_len > total_size: + raise RuntimeError( + "Exceeds total size: " + f"offset={offset}, req.extend_input_len={req.extend_input_len}, " + f"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: + if seq_len - pre_len != req.extend_input_len: + raise RuntimeError( + f"seq_len={seq_len}, pre_len={pre_len}, " + f"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.seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.int64, pin_memory=True) + 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 + + # 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]) + alloced_len = len(req.fill_ids) - 1 + self.req_to_token_pool.set_req_pool_info( + req.req_pool_idx, + ReqToTokenPoolInfo( + alloced_len, + alloced_len, + self.req_to_token_pool.req_to_token[ + req.req_pool_idx, :alloced_len + ].clone(), + ), + ) + # Cache the request in tree_cache with full sequence + self.tree_cache.cache_unfinished_req(req) + if req.grammar is not None: + req.grammar.accept_token(req.output_ids[-1]) + req.grammar.finished = req.finished() + self.output_ids = torch.tensor( + self.output_ids, device=self.device, dtype=torch.int32 + ) + + # 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(): + + self.prealloc_for_draft_decode(is_disaggregation_decode=True) + b, topk = len(self.reqs), server_args.speculative_eagle_topk + if topk != 1: + raise ValueError("Tree attention is abandoned for now") + last_verified_ids, token_list = self.output_ids, [] + + for _ in range(server_args.speculative_num_steps): + topk_index = torch.arange( + b * topk, device=self.device, dtype=torch.int32 + ) + topk_index = topk_index.reshape(b, topk) # shape: (b, topk) + token_list.append(topk_index) + + # local import to avoid circular importx + from tokenspeed.runtime.spec_decode.eagle import EagleDraftOutput + + # use draft output to create verify input next + spec_info = EagleDraftOutput( + last_verified_ids=last_verified_ids, + token_list=torch.cat(token_list, dim=-1), + ) + self.spec_info = spec_info diff --git a/python/tokenspeed/runtime/pd/epd/conn.py b/python/tokenspeed/runtime/pd/epd/conn.py new file mode 100644 index 0000000..fc8908e --- /dev/null +++ b/python/tokenspeed/runtime/pd/epd/conn.py @@ -0,0 +1,83 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Embedding-role transport substrate for EPD (encode -> prefill). + +The mechanical plumbing (the Mooncake engine + ZMQ control socket + room-keyed +status FSM, and the bootstrap rendezvous server) is the role-neutral shared base +in :mod:`tokenspeed.runtime.pd.base`; this module just layers the embedding role's +buffer args onto it. The embedding-specific *semantics* -- the wire frames and +the senders/receivers -- live in :mod:`embedding_transfer`; nothing here carries +any KV/MLA/metrics state. +""" + +from __future__ import annotations + +from tokenspeed.runtime.pd.base.bootstrap import DisaggBootstrapServerBase +from tokenspeed.runtime.pd.base.manager import DisaggManagerBase +from tokenspeed.runtime.pd.base.mooncake_engine import ( + MooncakeTransferEngine, +) +from tokenspeed.runtime.pd.epd.entities import EmbeddingManagerArgs +from tokenspeed.runtime.pd.utils import DisaggregationMode +from tokenspeed.runtime.utils.network import get_local_ip_by_remote + + +class MooncakeEmbeddingManagerBase(DisaggManagerBase): + """Embedding (encode->prefill) manager: the shared engine/socket/status FSM + (:class:`tokenspeed.runtime.pd.base.manager.DisaggManagerBase`) plus the + embedding buffer args. + Carries none of the KV manager's KV/MLA/metrics fields. + """ + + def __init__( + self, + args: EmbeddingManagerArgs, + embedding_args, # EmbeddingArgs: embedding buffer ptrs + gpu/ib device + disaggregation_mode: DisaggregationMode, + ): + self.args = args + self.embedding_args = embedding_args + self.disaggregation_mode = disaggregation_mode + self.bootstrap_port = args.bootstrap_port + self.world_size = args.tp_size + self.dp_size = 1 + # embedding_args must be set above before super().__init__, which calls + # register_buffer_to_engine. + engine = MooncakeTransferEngine( + hostname=get_local_ip_by_remote(), + gpu_id=embedding_args.gpu_id, + ib_device=embedding_args.ib_device, + ) + super().__init__(engine=engine) + + def register_buffer_to_engine(self): + ea = self.embedding_args + if ea.embedding_data_ptr: + self.engine.register(ea.embedding_data_ptr, ea.embedding_data_len) + if ea.deepstack_data_ptr: + self.engine.register(ea.deepstack_data_ptr, ea.deepstack_data_len) + + +class MooncakeEmbeddingBootstrapServer(DisaggBootstrapServerBase): + """Embedding bootstrap rendezvous: the shared server with no extra fields -- + the encode->prefill handshake needs only the base ip/port + parallel-size + sync, so there is nothing to layer onto :class:`DisaggBootstrapServerBase`. + """ diff --git a/python/tokenspeed/runtime/pd/epd/embedding_transfer.py b/python/tokenspeed/runtime/pd/epd/embedding_transfer.py new file mode 100644 index 0000000..7296780 --- /dev/null +++ b/python/tokenspeed/runtime/pd/epd/embedding_transfer.py @@ -0,0 +1,722 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Mooncake transport for EPD encode->prefill image-embedding transfer. + +The encode (vision-tower-only) server is the data source: it calls +``batch_transfer_sync`` to ship a contiguous ``[num_tokens, hidden]`` tensor +per item over the same Mooncake RDMA engine the KV path uses. The prefill side +is the receiver that pre-registers buffers and registers with the bootstrap +server (direction reversed relative to prefill->decode). + +Wire-frame dataclasses live in :mod:`tokenspeed.runtime.pd.epd.entities`: they +are pure data plus codecs and import no torch, so the protocol contract stays +importable and unit-testable on CPU. +""" + +from __future__ import annotations + +import concurrent.futures +import os +import threading +import time +from typing import TYPE_CHECKING + +import requests +import zmq + +from tokenspeed.runtime.pd.base.status import TransferPoll +from tokenspeed.runtime.pd.epd.conn import ( + MooncakeEmbeddingManagerBase, +) +from tokenspeed.runtime.pd.epd.entities import ( + REGISTER_ROOM_SENTINEL, + EmbeddingArgs, + EmbeddingArgsRegisterInfo, + EmbeddingChunk, + EmbeddingManagerArgs, + EmbeddingTransferError, + EmbeddingTransferInfo, +) +from tokenspeed.runtime.pd.utils import DisaggregationMode, FastQueue +from tokenspeed.runtime.utils import get_colorful_logger +from tokenspeed.runtime.utils.env import envs +from tokenspeed.runtime.utils.network import get_free_port, get_local_ip_by_remote + +logger = get_colorful_logger(__name__) + +if TYPE_CHECKING: + import torch + + +def _route_get(bootstrap_addr: str, engine_rank: int, target_dp_group: int): + """GET the bootstrap server's /route endpoint; ``None`` on any failure. + + Must never raise: callers in the receiver ``__init__`` treat ``None`` as a + per-room failure, but an uncaught exception would escape the prefill + scheduler thread and take the whole engine down via SIGUSR1. + """ + url = ( + f"http://{bootstrap_addr}/route?" + f"engine_rank={engine_rank}&target_dp_group={target_dp_group}" + ) + try: + resp = requests.get(url, timeout=5) + except Exception as e: # noqa: BLE001 -- any transport failure -> per-room fail + logger.error("EPD bootstrap /route fetch failed (%s): %s", url, e) + return None + if resp.status_code == 200: + return resp.json() + return None + + +def _b(value: object) -> bytes: + return str(value).encode("ascii") + + +def validate_fanout_frames( + infos: list[EmbeddingTransferInfo], chunk: "EmbeddingChunk" +) -> str | None: + """Pre-send contract check over a room's full fanout set; None when valid. + + Per frame: hidden/dtype must match the chunk; ``span`` (every frame carries + the image's full row count) must equal the chunk's token count (the + unchecked RDMA write silently truncates on divergence); the shard must stay + inside the chunk's rows; and deepstack presence must agree on both sides. + The set must then be either ALL identity (each frame covers the full span -- + full-copy broadcast) or a proper shard set whose non-empty shards, sorted by + ``row_start``, tile a contiguous range disjointly (an encode rank under + encode_tp>1 serves a contiguous BLOCK of the global shards). A gap/overlap + means the two sides' shard math diverged, so the room must fail loud. + """ + for info in infos: + if not (info.hidden == chunk.hidden and info.dtype == chunk.dtype): + return "embedding shape/dtype contract violated" + if info.span != chunk.n_tokens: + return ( + f"image token-count contract violated: receiver expects " + f"{info.span} rows, encode has {chunk.n_tokens} (G2)" + ) + if ( + info.row_start < 0 + or info.n_tokens < 0 + or info.row_start + info.n_tokens > chunk.n_tokens + ): + return ( + f"embedding shard out of range: rows [{info.row_start}, " + f"{info.row_start + info.n_tokens}) of {chunk.n_tokens}" + ) + if info.has_deepstack and not chunk.deepstack_width: + return ( + "receiver expects deepstack but the chunk carries none " + "(encode-side cache hit without its deepstack half?)" + ) + if chunk.deepstack_width and not info.has_deepstack: + return "chunk carries deepstack but the receiver did not allocate for it" + if all(i.row_start == 0 and i.n_tokens == chunk.n_tokens for i in infos): + return None # identity (full-copy) mode + shards = sorted((i.row_start, i.n_tokens) for i in infos if i.n_tokens > 0) + for (start, count), (next_start, _next_count) in zip(shards, shards[1:]): + if start + count != next_start: + return ( + "embedding shard frames do not tile contiguously: " + f"[{start}, {start + count}) then [{next_start}, ...)" + ) + return None + + +def shard_payload( + chunk: "EmbeddingChunk", info: EmbeddingTransferInfo +) -> tuple[int, int, int, int]: + """Source pointer/length math for one receiver's row shard of a chunk. + + Returns ``(src_ptr, nbytes, src_deepstack_ptr, deepstack_nbytes)`` for the + rows ``[info.row_start, info.row_start + info.n_tokens)``. Identity frames + reproduce the whole-chunk payload exactly. Row strides are derived from the + chunk's own byte counts so dtype size never needs decoding here. + """ + row_bytes = chunk.nbytes // chunk.n_tokens if chunk.n_tokens else 0 + src = chunk.src_embedding_ptr + info.row_start * row_bytes + nbytes = info.n_tokens * row_bytes + deep_src = deep_nbytes = 0 + if chunk.deepstack_width and chunk.deepstack_nbytes: + deep_row_bytes = ( + chunk.deepstack_nbytes // chunk.n_tokens if chunk.n_tokens else 0 + ) + deep_src = chunk.src_deepstack_ptr + info.row_start * deep_row_bytes + deep_nbytes = info.n_tokens * deep_row_bytes + return src, nbytes, deep_src, deep_nbytes + + +class MooncakeEmbeddingSender: + """Encode-side per-request sender for one bootstrap room. + + ``poll`` / ``clear`` / ``failure_exception`` are status-only (they touch + only the manager's status maps). ``send`` queues one contiguous embedding + tensor described by scalar fields, keeping this class free of any torch + dependency. + """ + + def __init__(self, mgr, bootstrap_addr: str, bootstrap_room: int): + self.mgr = mgr # MooncakeEmbeddingManagerEncode + self.bootstrap_server_url = bootstrap_addr + self.bootstrap_room = bootstrap_room + self.mgr.update_status(bootstrap_room, TransferPoll.Bootstrapping) + self.conclude_state = None + + def send( + self, + *, + src_embedding_ptr: int, + n_tokens: int, + hidden: int, + dtype: str, + nbytes: int, + src_deepstack_ptr: int = 0, + deepstack_width: int = 0, + deepstack_nbytes: int = 0, + copy_event: torch.cuda.Event | None = None, + ) -> None: + chunk = EmbeddingChunk( + room=self.bootstrap_room, + src_embedding_ptr=src_embedding_ptr, + n_tokens=n_tokens, + hidden=hidden, + dtype=dtype, + nbytes=nbytes, + src_deepstack_ptr=src_deepstack_ptr, + deepstack_width=deepstack_width, + deepstack_nbytes=deepstack_nbytes, + copy_event=copy_event, + ) + self.mgr.add_transfer_request(self.bootstrap_room, chunk) + + def poll(self) -> int: + # No Bootstrapping timeout here: the never-registered-receiver case is + # handled by the manager's _park_reaper (it fails rooms whose parked chunk + # outlives bootstrap_time_out), so this only reports terminal status. + if self.conclude_state is None: + status = self.mgr.check_status(self.bootstrap_room) + if status in (TransferPoll.Success, TransferPoll.Failed): + self.conclude_state = status + return status + return self.conclude_state + + def clear(self) -> None: + if self.bootstrap_room in self.mgr.request_status: + self.mgr.request_status.pop(self.bootstrap_room) + + def failure_exception(self): + if self.conclude_state is None: + self.conclude_state = TransferPoll.Failed + self.clear() + with self.mgr.failure_lock: + failure_reason = self.mgr.failure_records.pop( + self.bootstrap_room, "Failed due to an unknown reason from another rank" + ) + raise EmbeddingTransferError( + self.bootstrap_room, failure_reason, self.bootstrap_server_url + ) + + +class MooncakeEmbeddingManagerEncode(MooncakeEmbeddingManagerBase): + """Encode-side (data source) manager: registers itself to the bootstrap + server so the prefill receiver can find it, listens for the receiver's + registration + per-request pre-alloc frames, and a worker thread issues the + one-sided Mooncake write of the embedding. + """ + + def __init__(self, args: EmbeddingManagerArgs, embedding_args: EmbeddingArgs): + super().__init__(args, embedding_args, DisaggregationMode.ENCODE) + # room -> {receiver_session -> EmbeddingTransferInfo} + self.transfer_infos: dict[int, dict[str, EmbeddingTransferInfo]] = {} + # Bootstrap registration-wait timeout (default 120s). + self.bootstrap_time_out = envs.TOKENSPEED_DISAGGREGATION_BOOTSTRAP_TIMEOUT.get() + # room -> [(deadline, chunk), ...] chunks popped before the receiver + # registered; flushed by the bootstrap thread once its info arrives. + self._pending: dict[int, list[tuple]] = {} + self._pending_lock = threading.Lock() + # zmq sockets are not thread-safe and status pushes come from any + # fanout-pool thread: keep one PUSH socket per (thread, endpoint). + self._status_tls = threading.local() + self._start_bootstrap_thread() + self._register_to_bootstrap() + # K queues sharded by ROOM: a single consumer per room preserves the + # park/flush/straggler-drop ordering the parking machinery relies on. + # Each queue's worker uses a pool to issue one room's per-receiver writes + # concurrently (each prefill rank is a distinct Mooncake session, so a + # single batch call cannot span them). + cpu_count = os.cpu_count() or 8 + pool_size = envs.TOKENSPEED_DISAGGREGATION_THREAD_POOL_SIZE.get_set_value_or( + min(max(4, int(0.75 * cpu_count) // 8), 12) + ) + queue_count = envs.TOKENSPEED_DISAGGREGATION_QUEUE_SIZE.get() + assert pool_size >= queue_count, ( + f"TOKENSPEED_DISAGGREGATION_THREAD_POOL_SIZE={pool_size} must be >= " + f"TOKENSPEED_DISAGGREGATION_QUEUE_SIZE={queue_count}" + ) + self._queues: list[FastQueue] = [FastQueue() for _ in range(queue_count)] + self._executors = [ + concurrent.futures.ThreadPoolExecutor(max(1, pool_size // queue_count)) + for _ in range(queue_count) + ] + for queue, executor in zip(self._queues, self._executors): + threading.Thread( + target=self._transfer_worker, args=(queue, executor), daemon=True + ).start() + threading.Thread(target=self._park_reaper, daemon=True).start() + + def _start_bootstrap_thread(self): + self.rank_port = get_free_port() + self.server_socket.bind(f"tcp://{get_local_ip_by_remote()}:{self.rank_port}") + + def loop(): + while True: + msg = self.server_socket.recv_multipart() + # A malformed frame must not kill this daemon thread: a dead + # listener drops every later registration, parking all chunks + # until the reaper fails them. Log and continue. + try: + if msg[0].decode("ascii") == REGISTER_ROOM_SENTINEL: + # registration frame: consumed, no per-session state needed + pass + else: + info = EmbeddingTransferInfo.from_zmq(msg) + self.transfer_infos.setdefault(info.room, {})[ + info.mooncake_session_id + ] = info + if ( + len(self.transfer_infos[info.room]) + >= info.required_dst_info_num + ): + self.update_status(info.room, TransferPoll.Bootstrapped) + # Re-drive any chunk parked before this room registered. + self._flush_pending(info.room) + except Exception as exc: # noqa: BLE001 + logger.error( + "dropping malformed embedding bootstrap frame " + "(%d parts): %s", + len(msg), + exc, + ) + + threading.Thread(target=loop, daemon=True).start() + + def _register_to_bootstrap(self): + ip = get_local_ip_by_remote() + bootstrap_host = self.args.bootstrap_host or ip + payload = { + "role": "Prefill", # bootstrap-server role string for "discoverable data source" + "world_size": self.world_size, + "dp_size": self.dp_size, + "rank_ip": ip, + "rank_port": self.rank_port, + "engine_rank": self.embedding_args.engine_rank, + } + url = f"http://{bootstrap_host}:{self.bootstrap_port}/route" + # The bootstrap HTTP server starts concurrently with this call, so the + # first PUTs can race it and hit connection-refused. A dropped + # registration leaves /route's parallel-info null and the encode silently + # serves no embeddings, so retry until the server accepts it. + last_err = None + for _ in range(60): + try: + resp = requests.put(url, json=payload, timeout=5) + if resp.ok: + return + last_err = f"status {resp.status_code}" + except Exception as e: # noqa: BLE001 + last_err = e + time.sleep(0.5) + logger.error( + "encode failed to register to bootstrap server after retries: %s", last_err + ) + + def add_transfer_request(self, room: int, chunk: EmbeddingChunk) -> None: + if ( + room not in self.request_status + or self.check_status(room) == TransferPoll.Failed + ): + return + self._queue_for(room).put(chunk) + + def _queue_for(self, room: int) -> FastQueue: + # Room-affinity sharding: one consumer per room keeps the park / + # _flush_pending re-enqueue / Success-straggler-drop sequence + # single-threaded per room. + return self._queues[room % len(self._queues)] + + def _transfer_worker( + self, queue: FastQueue, executor: concurrent.futures.ThreadPoolExecutor + ): + while True: + chunk: EmbeddingChunk = queue.get() + # Already concluded: a re-enqueued straggler copy from _flush_pending + # (which re-puts a parked chunk on each receiver registration). The + # room was already sent + popped, so drop it. + if self.request_status.get(chunk.room) == TransferPoll.Success: + continue + # Snapshot via list(): the bootstrap listener inserts into this inner + # dict without _pending_lock, so next(iter(...)) could raise "dict changed + # size during iteration". list() is GIL-atomic. + info_vals = list(self.transfer_infos.get(chunk.room, {}).values()) + # required_dst_info_num (= fanout) is identical across the N frames for a + # room, so read it off any registered frame; None until one registers. + need = info_vals[0].required_dst_info_num if info_vals else None + if (need is None or len(info_vals) < need) and self.request_status.get( + chunk.room + ) != TransferPoll.Failed: + # Fewer than N receivers registered (1->N broadcast not yet complete): + # park (don't drop) until the full set arrives. _flush_pending re-drives + # this chunk on every new registration; the reaper fails it on timeout. + parked = True + with self._pending_lock: + info_vals = list(self.transfer_infos.get(chunk.room, {}).values()) + need = info_vals[0].required_dst_info_num if info_vals else None + if info_vals and need is not None and len(info_vals) >= need: + # Full set landed inside the park window: send below instead. + parked = False + else: + self._pending.setdefault(chunk.room, []).append( + (time.time() + self.bootstrap_time_out, chunk) + ) + if parked: + continue + # All N receivers registered: serve every one concurrently (full copy + # in identity mode, its row shard otherwise) -- each is a distinct + # Mooncake session, so concurrency must come from threads. + infos = list(self.transfer_infos.get(chunk.room, {}).values()) + err = validate_fanout_frames(infos, chunk) + if err is not None: + self._fail_room(chunk.room, err, infos) + continue + # Wait the ring copy's completion HERE on the daemon (not the + # encode-loop thread): the one-sided RDMA reads in _send touch GPU + # memory off any CUDA stream, so they must not precede the device copy + # that filled the slot (ViT->send corruption hazard). One wait per + # chunk covers all N receivers (same slot); event.synchronize() + # releases the GIL while waiting. + if chunk.copy_event is not None: + chunk.copy_event.synchronize() + futures = [executor.submit(self._send, info, chunk) for info in infos] + failed = False + for info, future in zip(infos, futures): + ret = future.result() # _send returns <0 on error, never raises + if ret != 0: + self.record_failure(chunk.room, f"mooncake transfer ret={ret}") + failed = True + else: + # Per-receiver completion sync (each of the N prefill ranks gets + # exactly one, matching its required_response_num == 1). + self._sync_status(info, TransferPoll.Success) + if failed: + # Push Failed to EVERY receiver (idempotent on ones already + # Success'd: the rank-synced admission MIN aborts the request + # everywhere anyway). + self._fail_room(chunk.room, None, infos) + continue + # All N receivers served: mark the room Success ONCE and pop only now, + # so a partial fanout can never conclude/pop early. + self.update_status(chunk.room, TransferPoll.Success) + self.transfer_infos.pop(chunk.room, None) + + def _fail_room( + self, + room: int, + reason: str | None, + infos: list[EmbeddingTransferInfo], + ) -> None: + if reason is not None: + self.record_failure(room, reason) + self.update_status(room, TransferPoll.Failed) + for info in infos: + self._sync_status(info, TransferPoll.Failed) + + def is_parked(self, room: int) -> bool: + """Whether ``room`` still has a chunk parked awaiting receiver + registration. Public probe for the executor's ring-slot lease: a parked + chunk holds its slot's pointer for re-send, so the slot is not reusable + until the room unparks.""" + with self._pending_lock: + return room in self._pending + + def fail_room(self, room: int, reason: str | None) -> None: + """Conclude ``room`` Failed and push Failed to all of its registered + receivers. Public seam for the encode executor, which must not reach into + ``transfer_infos`` or the status FSM directly.""" + infos = list(self.transfer_infos.get(room, {}).values()) + self._fail_room(room, reason, infos) + + def _flush_pending(self, room: int) -> None: + # Re-enqueue this room's parked chunks (in order) now that it registered. + # Same room -> same queue, so the re-driven copy is consumed by the same + # single worker that parked it. + with self._pending_lock: + parked = self._pending.pop(room, []) + for _deadline, chunk in parked: + self._queue_for(room).put(chunk) + + def _park_reaper(self): + # Fail rooms whose parked chunk outlived bootstrap_time_out (the + # receiver never registered), so an aborted request never hangs. + while True: + time.sleep(1.0) + now = time.time() + with self._pending_lock: + expired = [ + room + for room, parked in self._pending.items() + if parked and now >= parked[0][0] + ] + for room in expired: + self._pending.pop(room, None) + for room in expired: + self.record_failure( + room, f"receiver never registered within {self.bootstrap_time_out}s" + ) + self.update_status(room, TransferPoll.Failed) + + def _send(self, info: EmbeddingTransferInfo, chunk: EmbeddingChunk) -> int: + src, nbytes, deep_src, deep_nbytes = shard_payload(chunk, info) + if nbytes == 0: + # Zero-row shard (image span < shard count): nothing to write, but + # the frame already served as this receiver's registration and it + # still gets its Success push so its poll() completes. + return 0 + bufs = [src] + dsts = [info.dst_embedding_ptr] + lens = [nbytes] + if deep_nbytes and info.dst_deepstack_ptr: + bufs.append(deep_src) + dsts.append(info.dst_deepstack_ptr) + lens.append(deep_nbytes) + # Always go through the BATCH transfer API, even for a single buffer: its + # binding releases the GIL for the one-sided RDMA write + # (gil_scoped_release), while singular transfer_sync_write does not + # reliably, pinning the GIL for the whole write and freezing this daemon's + # loop + other daemons. A 1-element batch is byte-for-byte the same + # transfer. + return self.engine.batch_transfer_sync( + info.mooncake_session_id, bufs, dsts, lens + ) + + def _sync_status(self, info: EmbeddingTransferInfo, status: int) -> None: + # One socket per (thread, endpoint): pushes come from fanout-pool threads + # and zmq sockets are not thread-safe. + socks = getattr(self._status_tls, "socks", None) + if socks is None: + socks = self._status_tls.socks = {} + endpoint = f"tcp://{info.endpoint}:{info.dst_port}" + sock = socks.get(endpoint) + if sock is None: + sock = zmq.Context.instance().socket(zmq.PUSH) + sock.connect(endpoint) + socks[endpoint] = sock + sock.send_multipart( + [_b(info.room), _b(int(status)), _b(self.embedding_args.engine_rank)] + ) + + +class MooncakeEmbeddingManagerPrefill(MooncakeEmbeddingManagerBase): + """Prefill-side (data sink) manager: holds the discovery caches and a thread + that consumes the encode side's per-request completion-status frames, marking + the room Success once all expected responses arrive. + """ + + def __init__(self, args: EmbeddingManagerArgs, embedding_args: EmbeddingArgs): + super().__init__(args, embedding_args, DisaggregationMode.PREFILL) + self.required_response_num: dict[int, int] = {} + self.response_tracker: dict[int, set] = {} + self.connection_pool: dict[str, list] = {} + self.prefill_parallel_info: dict[str, dict] = {} + self._start_status_thread() + + def _start_status_thread(self): + self.rank_port = get_free_port() + self.server_socket.bind(f"tcp://{get_local_ip_by_remote()}:{self.rank_port}") + + def loop(): + while True: + parts = self.server_socket.recv_multipart() + room = int(parts[0].decode("ascii")) + status = int(parts[1].decode("ascii")) + rank = int(parts[2].decode("ascii")) + if status == TransferPoll.Success and room in self.request_status: + self.response_tracker.setdefault(room, set()).add(rank) + if len( + self.response_tracker[room] + ) >= self.required_response_num.get(room, 1): + self.update_status(room, TransferPoll.Success) + elif status == TransferPoll.Failed: + self.record_failure(room, "encode failed to send embedding") + self.update_status(room, TransferPoll.Failed) + + threading.Thread(target=loop, daemon=True).start() + + +class MooncakeEmbeddingReceiver: + """Prefill-side per-request receiver: discovers the encode endpoint via the + bootstrap server, registers its receive buffer, and on ``pre_alloc`` tells + the encode side where/how big to write this request's embedding. 1->N + broadcast (prefill_tp a multiple of encode_tp): contiguous blocks of prefill + ranks pair one encode rank; encode_tp=1 -> all prefill ranks pair encode + rank 0 and receive the same TP-gathered embedding. + """ + + def __init__( + self, + mgr: MooncakeEmbeddingManagerPrefill, + bootstrap_addr: str, + bootstrap_room: int, + ): + self.mgr = mgr + self.bootstrap_addr = bootstrap_addr + self.bootstrap_room = bootstrap_room + self.session_id = mgr.get_session_id() + mgr.update_status(bootstrap_room, TransferPoll.Bootstrapping) + + pinfo = mgr.prefill_parallel_info.get(bootstrap_addr) + if pinfo is None or pinfo.get("prefill_tp_size") is None: + # SINGLE-attempt fail-fast: this runs in the prefill scheduler thread, + # so a retry-loop would stall the whole TP group on a slow/unreachable + # encode bootstrap. /route can answer 200 with un-populated + # parallel-info while the encode worker is still registering (startup + # race); fail this room fast and let the client retry once the encode + # has registered. A partial dict is never cached (it would crash later + # requests on int(None)). _route_get bounds its own GET at timeout=5, + # so this attempt cannot block beyond that. + cand = _route_get(bootstrap_addr, -1, -1) + if cand is None or cand.get("prefill_tp_size") is None: + mgr.record_failure( + bootstrap_room, "no (complete) parallel info from bootstrap" + ) + mgr.update_status(bootstrap_room, TransferPoll.Failed) + return + pinfo = cand + mgr.prefill_parallel_info[bootstrap_addr] = pinfo + + encode_tp = int(pinfo["prefill_tp_size"]) + encode_dp = int(pinfo["prefill_dp_size"]) + local_tp = mgr.world_size // mgr.dp_size + # prefill_tp must be a whole multiple of encode_tp. The vision tower + # output is TP-gathered (identical on every encode rank) and every prefill + # rank needs the full embedding, so encode_tp=1 -> prefill_tp=N is a 1->N + # broadcast. Contiguous blocks of `fanout` prefill ranks share one encode + # rank. + assert local_tp % encode_tp == 0, ( + f"EPD requires prefill_tp to be a multiple of encode_tp " + f"(encode_tp={encode_tp}, prefill_tp={local_tp})" + ) + fanout = local_tp // encode_tp # prefill ranks served by one encode rank + # Drives the ENCODE-side gate: how many prefill ranks register with this + # request's encode rank before it may send + conclude (1->N broadcast). + self.required_dst_info_num = fanout + # This prefill rank still pulls from exactly ONE encode rank, so it expects + # a single completion sync. The N-way fan-out is purely the encode side's + # concern; bumping this to N would hang every prefill rank. + mgr.required_response_num[bootstrap_room] = 1 + + # Which encode rank this prefill rank pulls from: contiguous grouping. + # encode_tp=1 -> fanout=local_tp -> my_encode_rank == 0 for all prefill + # ranks. Use `// fanout`, NOT `% encode_tp`: both are 0 for encode_tp=1, + # but only `// fanout` groups ranks correctly for encode_tp>1. + my_tp_rank = mgr.embedding_args.engine_rank % local_tp + my_encode_rank = my_tp_rank // fanout + target_dp_group = bootstrap_room % encode_dp + key = f"{bootstrap_addr}_{target_dp_group}_{my_encode_rank}" + if key not in mgr.connection_pool: + info = _route_get(bootstrap_addr, my_encode_rank, target_dp_group) + if info is None: + mgr.record_failure(bootstrap_room, "no encode rank info from bootstrap") + mgr.update_status(bootstrap_room, TransferPoll.Failed) + return + self.encode_infos = [info] + mgr.connection_pool[key] = self.encode_infos + self._register_args() + else: + self.encode_infos = mgr.connection_pool[key] + mgr.update_status(bootstrap_room, TransferPoll.Bootstrapped) + + def _register_args(self): + ea = self.mgr.embedding_args + reg = EmbeddingArgsRegisterInfo( + room=REGISTER_ROOM_SENTINEL, + endpoint=get_local_ip_by_remote(), + dst_port=self.mgr.rank_port, + mooncake_session_id=self.session_id, + dst_embedding_ptr=ea.embedding_data_ptr, + dst_deepstack_ptr=ea.deepstack_data_ptr, + ) + for info in self.encode_infos: + sock = self.mgr._connect(f"tcp://{info['rank_ip']}:{info['rank_port']}") + sock.send_multipart(reg.to_zmq()) + + def pre_alloc( + self, + *, + dst_embedding_ptr: int, + n_tokens: int, + hidden: int, + dtype: str, + dst_deepstack_ptr: int = 0, + has_deepstack: bool = False, + row_start: int = 0, + span: int = 0, + ) -> None: + """Tell the encode side where/how big to write. In shard mode the caller + passes this rank's row sub-range: ``row_start`` within the image and + ``n_tokens`` = the SHARD's row count, both dst pointers already offset to + the shard's first row. ``span`` is the image's FULL row count (pass it in + identity mode too: it is the encode side's token-count tripwire). A + ``n_tokens == 0`` frame is still sent: it doubles as this receiver's + registration heartbeat (the encode-side fanout gate counts frames, not + bytes).""" + for info in self.encode_infos: + ti = EmbeddingTransferInfo( + room=self.bootstrap_room, + endpoint=get_local_ip_by_remote(), + dst_port=self.mgr.rank_port, + mooncake_session_id=self.session_id, + dst_embedding_ptr=dst_embedding_ptr, + dst_deepstack_ptr=dst_deepstack_ptr, + n_tokens=n_tokens, + hidden=hidden, + dtype=dtype, + has_deepstack=has_deepstack, + required_dst_info_num=self.required_dst_info_num, + row_start=row_start, + span=span, + ) + sock = self.mgr._connect(f"tcp://{info['rank_ip']}:{info['rank_port']}") + sock.send_multipart(ti.to_zmq()) + + def poll(self) -> int: + return self.mgr.check_status(self.bootstrap_room) + + def clear(self) -> None: + """Drop this room's bookkeeping from the (singleton) prefill manager on a + terminal receive, else every request leaks its status/tracker entries. + Mirrors the KV receiver; called post-terminal so it can't race check_status.""" + room = self.bootstrap_room + self.mgr.request_status.pop(room, None) + self.mgr.required_response_num.pop(room, None) + self.mgr.response_tracker.pop(room, None) + with self.mgr.failure_lock: + self.mgr.failure_records.pop(room, None) diff --git a/python/tokenspeed/runtime/pd/epd/encode_executor.py b/python/tokenspeed/runtime/pd/epd/encode_executor.py new file mode 100644 index 0000000..f928b49 --- /dev/null +++ b/python/tokenspeed/runtime/pd/epd/encode_executor.py @@ -0,0 +1,372 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""EPD encode-worker execution: run the vision tower, scatter its output back +onto each item, and hand the contiguous embeddings to the Mooncake sender. +""" + +from __future__ import annotations + +import logging + +import torch + +from tokenspeed.runtime.multimodal.embedder import _item_token_count +from tokenspeed.runtime.multimodal.inputs import Modality, MultimodalDataItem +from tokenspeed.runtime.pd.base.status import TransferPoll +from tokenspeed.runtime.pd.epd.embedding_transfer import ( + MooncakeEmbeddingSender, +) +from tokenspeed.runtime.utils.env import envs + +logger = logging.getLogger(__name__) + + +def assign_encoded_embeddings( + items: list[MultimodalDataItem], + output: torch.Tensor, + model, +) -> None: + """Scatter a packed vision-tower output onto each item, in place. + + ``output`` is the tower's ``[sum_tokens, width]`` result for ``items`` in + order (``width = hidden`` for plain models, or ``hidden * (1 + n_deepstack)`` + for deepstack models like Qwen3.5). Each item's row span is its post-merge + token count (``_item_token_count``); the rows are split accordingly and, + for deepstack models, column-split via ``model.separate_deepstack_embeds`` + into the main ``[N, hidden]`` and deepstack ``[N, hidden * n_deepstack]`` + halves. Results are made contiguous because a TP-gathered tower output may + not be, and the transfer ships raw row-major bytes. + + Sets ``item.encoded`` (and ``item.encoded_deepstack`` when the model emits + deepstack, else ``None``), which is exactly the ``skip-ViT`` form the + prefill-side VisionEmbedder consumes. + """ + output = output.reshape(-1, output.shape[-1]) + per_item_tokens = [_item_token_count(item) for item in items] + total = sum(per_item_tokens) + if output.shape[0] != total: + raise ValueError( + f"vision-tower output has {output.shape[0]} rows but items sum to " + f"{total} post-merge tokens; check the token-count / grid contract" + ) + + has_deepstack = getattr(model, "num_deepstack_embeddings", 0) > 0 + per_item_embeds = torch.split(output, per_item_tokens, dim=0) + for item, emb in zip(items, per_item_embeds): + if has_deepstack: + main, deep = model.separate_deepstack_embeds(emb) + item.encoded = main.contiguous() + item.encoded_deepstack = deep.contiguous() + else: + item.encoded = emb.contiguous() + item.encoded_deepstack = None + + +class DisaggEncodeExecutor: + """Drives one encode worker: run the vision tower on a batch of items, then + ship each item's embedding to its prefill peer over Mooncake. + + Python orchestration: this is invoked by the encode loop, not the + C++ scheduler. ``execute`` groups items by modality, runs the tower once per + modality via the model's ``get_image_feature`` / ``get_video_feature``, + scatters the output onto ``item.encoded`` (see + :func:`assign_encoded_embeddings`), and queues a transfer per item through + the per-request :class:`MooncakeEmbeddingSender`. Each request must first be + ``register``-ed with its prefill peer's bootstrap (host, port, room). + """ + + def __init__( + self, + manager, + multimodal_model, + device, + *, + ring_slots: int = 64, + ring_bytes: int = 256 * 1024 * 1024, + ): + self.manager = manager + self.model = multimodal_model + self.device = device + self.senders = {} + # RDMA requires every transferred buffer to be a registered memory region, + # and mooncake rejects OVERLAPPING registrations -- registering each + # per-request ``item.encoded`` fails because the torch caching allocator + # packs freed-but-still-registered tensors so a grown region straddles + # others. Collapse every send through a fixed ring of pre-registered bounce + # buffers: each slot is registered once at a fixed size (never grows, never + # overlaps), and ``item.encoded`` is copied into a slot before its async + # send. ``ring_slots`` / ``ring_bytes`` are injectable for tests and + # env-tunable; total reservation is slots * slot_bytes PER ring (main, plus + # deepstack if present), so depth and per-slot bytes must be sized to the + # model and peak concurrency. + self._ring_slots = int( + envs.TOKENSPEED_EPD_ENCODE_RING_SLOTS.get_set_value_or(ring_slots) + ) + slot_mb = envs.TOKENSPEED_EPD_ENCODE_RING_SLOT_MB.get() + # Env override is in whole MiB; unset -> keep the exact ``ring_bytes`` arg. + self._ring_bytes = slot_mb * 1024 * 1024 if slot_mb else ring_bytes + self._main_ring = None # lazily allocated on first send (device live by then) + self._deep_ring = None + self._ring_idx = 0 + # Per-slot lease: the room whose send last staged into the slot. A slot is + # reusable only once that room's transfer is TERMINAL and not parked (a + # parked chunk holds the slot's pointer until bootstrap_time_out and is + # re-sent on late receiver registration; see _lease_slot), so a full ring + # DEFERS the send rather than overwriting an in-flight slot. + self._slot_rooms: list = [None] * self._ring_slots + # Sends whose ViT output is ready but could not lease a free ring slot + # (all slots still hold in-flight transfers). Retried non-blocking by + # drain_deferred() each loop tick (a busy-wait here would GIL-starve the + # daemon transfer-workers that free the slots and deadlock the loop). + self._deferred_sends: list = [] + + def register(self, request_id, bootstrap_host, bootstrap_port, bootstrap_room): + self.senders[request_id] = MooncakeEmbeddingSender( + self.manager, f"{bootstrap_host}:{bootstrap_port}", bootstrap_room + ) + + def _feature_fn(self, modality): + # IMAGE dispatches through the model's ``image_encoder`` seam, NOT + # ``get_image_feature`` directly: that seam is what the encoder CUDA-graph + # wrapper overrides (see _maybe_install_encoder_cudagraph). When the graph + # is disabled the model leaves ``image_encoder = get_image_feature`` (eager); + # VIDEO has no captured graph -> always eager. + if modality == Modality.IMAGE: + return self.model.image_encoder + if modality == Modality.VIDEO: + return self.model.get_video_feature + raise ValueError(f"unsupported modality for encode: {modality}") + + def execute(self, request_items: list[tuple[str, MultimodalDataItem]]) -> None: + by_modality = {} + for _, item in request_items: + by_modality.setdefault(item.modality, []).append(item) + with torch.inference_mode(): + for modality, items in by_modality.items(): + output = self._feature_fn(modality)(items) + assign_encoded_embeddings(items, output, self.model) + # Stage every embedding into its ring slot, then issue the async + # Mooncake sends. See _stage_and_send for the copy/RDMA + # overwrite-safety invariant (one CUDA event gates each transfer). + self._stage_and_send(request_items) + + def _ensure_rings(self) -> None: + """Lazily allocate + register the bounce-buffer ring (see ``__init__``).""" + if self._main_ring is not None: + return + self._main_ring = [ + torch.empty(self._ring_bytes, dtype=torch.uint8, device=self.device) + for _ in range(self._ring_slots) + ] + for buf in self._main_ring: + self.manager.engine.register(buf.data_ptr(), self._ring_bytes) + + def _copy_into(self, ring, slot: int, src) -> tuple[int, int]: + """Copy ``src``'s bytes into pre-registered ring ``slot``; return its + (device pointer, byte length). Fails loud if an embedding exceeds a slot.""" + nbytes = src.numel() * src.element_size() + if nbytes > self._ring_bytes: + raise RuntimeError( + f"EPD encode embedding {nbytes} B exceeds ring slot " + f"{self._ring_bytes} B; raise TOKENSPEED_EPD_ENCODE_RING_SLOT_MB " + "or the ring_bytes constructor argument" + ) + buf = ring[slot] + buf[:nbytes].view(src.dtype).copy_(src.reshape(-1)) + return buf.data_ptr(), nbytes + + def _lease_slot(self) -> "int | None": + """Return a reusable ring-slot index, or ``None`` if every slot still + holds an in-flight transfer. NON-BLOCKING: the caller DEFERS rather than + spinning (a busy-wait would GIL-starve the daemon transfer-workers that + mark rooms terminal and deadlock the single-threaded loop). + + A slot is reusable once the room it last staged is TERMINAL (Success, + Failed, or None=already reaped) AND no parked chunk still holds its + pointer -- the overwrite-safety invariant.""" + mgr = self.manager + n = self._ring_slots + for _ in range(n): + slot = self._ring_idx % n + self._ring_idx += 1 + room = self._slot_rooms[slot] + if room is None: + return slot + status = mgr.room_status(room) + if status is None or status in ( + TransferPoll.Success, + TransferPoll.Failed, + ): + if not mgr.is_parked(room): + return slot + return None + + def _stage_and_send(self, items: list[tuple[str, MultimodalDataItem]]) -> None: + """Lease a ring slot per item and ship it; items that cannot lease a free + slot (ring full) are DEFERRED for a later non-blocking retry rather than + blocking the loop. Stages every leased item then issues ONE stream sync + before the sends, so the one-sided RDMA reads never race the device-to- + device copies (the ViT->send corruption hazard).""" + self._ensure_rings() + staged = [] + for rid, item in items: + if rid not in self.senders: + # Sender reaped (its room concluded/failed) -- drop this stale + # deferred send instead of crashing on senders[rid]. + continue + slot = self._lease_slot() + if slot is None: + self._deferred_sends.append((rid, item)) + continue + try: + send_args = self._stage_item( + item, self.senders[rid].bootstrap_room, slot + ) + except Exception as e: + # A staging error (most plausibly _copy_into rejecting an embedding + # larger than a ring slot) must fail only THIS item's room, never + # raise out of the single-threaded encode loop into the engine's + # SIGUSR1 handler (which kills the whole worker and every other + # in-flight image). Covers the unguarded send_item() / + # drain_deferred() callers too; the leased slot returns to the ring + # once the room is Failed (see _lease_slot). + self._fail_staged_room(rid, e) + continue + staged.append((rid, send_args)) + if not staged: + return + # Record ONE CUDA event after all the ring copies above (they ran on the + # current stream inside _stage_item) and hand it to each transfer rather + # than host-syncing on this single encode-loop thread. The daemon + # transfer-worker waits the event before its one-sided RDMA read + # (embedding_transfer._transfer_worker), so the read never races the copy; + # _lease_slot keeps the slot until its room is terminal (Success only after + # the RDMA completes). + copy_event = None + if torch.cuda.is_available(): + copy_event = torch.cuda.Event() + copy_event.record() + for rid, send_args in staged: + self.senders[rid].send(copy_event=copy_event, **send_args) + + def drain_deferred(self) -> None: + """Retry deferred sends (ViT done, waiting for a free ring slot). Non- + blocking: items that still cannot lease a slot stay deferred. Driven once + per encode-loop tick; the loop yields the GIL between ticks so the daemon + transfer-workers can free slots for the next drain.""" + if not self._deferred_sends: + return + pending = self._deferred_sends + self._deferred_sends = [] + self._stage_and_send(pending) + + def has_deferred(self) -> bool: + return bool(self._deferred_sends) + + def _conclude_room_failed(self, room: int, exc: Exception) -> None: + """Push Failed to ``room``'s prefill receivers so they abort via the + rank-synced admission path, instead of the error escaping the encode loop + and SIGUSR1-ing the worker. The single seam every failure path goes + through; delegates to the manager's public ``fail_room`` rather than + reaching into its ``transfer_infos`` / status FSM.""" + self.manager.fail_room(room, str(exc)) + + def _fail_staged_room(self, rid: str, exc: Exception) -> None: + """Per-item staging failure (the unguarded ``send_item`` / + ``drain_deferred`` callers): conclude ``rid``'s room Failed.""" + sender = self.senders.get(rid) + if sender is None: + return + self._conclude_room_failed(sender.bootstrap_room, exc) + logger.error( + "encode staging failed for room %s: %s", sender.bootstrap_room, exc + ) + + def fail_rooms(self, request_ids, exc: Exception) -> int: + """Conclude every room owned by ``request_ids`` Failed; return the count. + The owning seam for a batch-level failure that fired before any send was + issued (ViT / assign_encoded_embeddings): the worker hands its batch's + request_ids and stays out of the sender/manager internals. Rooms are + de-duped (a multi-image request shares one room).""" + rooms = set() + for rid in request_ids: + sender = self.senders.get(rid) + if sender is not None: + rooms.add(sender.bootstrap_room) + for room in rooms: + self._conclude_room_failed(room, exc) + return len(rooms) + + def reap_concluded_senders(self, pending_request_ids) -> None: + """Drop per-request senders whose room reached a terminal transfer status + (the ``senders`` dict otherwise grows forever). Senders whose request_id + is still awaiting the tower (``pending_request_ids``) are kept -- their + send has not been queued. Only the sender is dropped; the manager's + terminal ``request_status`` tombstone stays (the transfer worker's + straggler-drop and the ring-slot lease both key on it).""" + for rid in list(self.senders): + if rid in pending_request_ids: + continue + room = self.senders[rid].bootstrap_room + if self.manager.room_status(room) in ( + TransferPoll.Success, + TransferPoll.Failed, + ): + self.senders.pop(rid, None) + + def _stage_item(self, item: MultimodalDataItem, room, slot: int) -> dict: + """Copy one item's embedding (and deepstack half, if any) into the leased + ring ``slot`` and return the scalar ``send`` kwargs. The copy runs on the + current stream; the CALLER must synchronize before handing these pointers + to the transfer engine, so the one-sided RDMA read never races the device- + to-device copy (same hazard class as the ViT->send race).""" + enc = item.encoded + self._slot_rooms[slot] = room + send_ptr, nbytes = self._copy_into(self._main_ring, slot, enc) + ds_ptr = ds_width = ds_nbytes = 0 + deep = item.encoded_deepstack + if deep is not None and deep.numel() > 0: + if self._deep_ring is None: + self._deep_ring = [ + torch.empty(self._ring_bytes, dtype=torch.uint8, device=self.device) + for _ in range(self._ring_slots) + ] + for buf in self._deep_ring: + self.manager.engine.register(buf.data_ptr(), self._ring_bytes) + ds_width = deep.shape[1] + ds_ptr, ds_nbytes = self._copy_into(self._deep_ring, slot, deep) + return dict( + src_embedding_ptr=send_ptr, + n_tokens=enc.shape[0], + hidden=enc.shape[1], + dtype=str(enc.dtype), + nbytes=nbytes, + src_deepstack_ptr=ds_ptr, + deepstack_width=ds_width, + deepstack_nbytes=ds_nbytes, + ) + + def send_item(self, request_id, item: MultimodalDataItem) -> None: + """Ship an already-encoded item (``item.encoded`` set) to its prefill peer. + Used by the encode loop for cache hits, which skip the tower but still + transfer. Routes through the same lease-or-defer path as ``execute`` so a + full ring defers (non-blocking) instead of stalling the loop.""" + self._stage_and_send([(request_id, item)]) diff --git a/python/tokenspeed/runtime/pd/epd/encode_loop.py b/python/tokenspeed/runtime/pd/epd/encode_loop.py new file mode 100644 index 0000000..3c3d82b --- /dev/null +++ b/python/tokenspeed/runtime/pd/epd/encode_loop.py @@ -0,0 +1,315 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""EPD encode-worker loop: a lightweight, LM-free scheduler subprocess. + +The encode role runs the vision tower and ships image embeddings to prefill +workers over Mooncake; it owns no KV cache and never runs the language model, so +it does NOT use the full :class:`EventLoop`/C++ scheduler (that machinery -- paged +KV, chunked prefill, retract, token budget -- is an impedance mismatch for a ViT). +Instead ``run_event_loop`` branches here when ``disaggregation_mode == "encode"``. + +The assembly is: load the model, stand up the Mooncake encode manager + bootstrap +server, the DisaggEncodeExecutor, and the EncodeWorker, reusing the SAME request +IPC the LM scheduler uses: a ZMQ PULL +on ``port_args.scheduler_input_ipc_name``. The smg grpc_servicer's TokenSpeedEncoder +handler sends an :class:`EncodeRequest` (pickled) over that channel; the loop +drains them into ``EncodeWorker.submit`` and runs ``step`` to encode + transfer. + +TP: the vision tower is TP-sharded, so all encode ranks run each batch in lockstep +-- rank 0 owns the gateway ZMQ and broadcasts every batch to the TP group. The +embedding transport is a 1->N broadcast: prefill_tp must be a multiple of +encode_tp. Data-parallel encode workers inside one server are not supported; +horizontal scale comes from multiple independent encode servers selected by the +gateway. +""" + +from __future__ import annotations + +import time + +import zmq + +from tokenspeed.runtime.cache.embedding_cache import ( + EmbeddingCache, + TieredEmbeddingCache, +) +from tokenspeed.runtime.pd.epd.encode_scheduler import EncodeScheduler +from tokenspeed.runtime.pd.epd.encode_worker import EncodeWorker +from tokenspeed.runtime.utils import get_colorful_logger, get_zmq_socket +from tokenspeed.runtime.utils.env import envs + +logger = get_colorful_logger(__name__) + + +# Vision-embedding cache capacity. L1 lives in GPU VRAM; the optional L2 lives in +# host DRAM and catches L1 evictions so duplicate images skip the tower even past +# the VRAM working set. Both are whole-MiB env overrides. L2 defaults to 0 +# (disabled): the host tier is opt-in. NOTE both knobs are PER ENCODE PROCESS (per +# TP rank): at encode TP>1, every co-located rank allocates its own L1+L2, so +# budget host DRAM as tp_size * EMBED_CACHE_DRAM_MB. +def _embedding_cache_bytes(env_field) -> int: + """Whole-MiB env field -> bytes. + + EnvField handles parsing and defaults; negative capacities are rejected here + with an env-named error. + """ + mb = env_field.get() + if mb < 0: + raise ValueError(f"{env_field.name} must be >= 0 MiB, got {mb}") + return mb * 1024 * 1024 + + +def _make_embedding_cache(l1_bytes: int, l2_bytes: int, device: str): + """Select the encode embedding cache: a plain single-tier VRAM + :class:`EmbeddingCache` by default, or a two-tier :class:`TieredEmbeddingCache` + (VRAM L1 + host-DRAM L2) when the L2 capacity is enabled (``l2_bytes > 0``).""" + if l2_bytes > 0: + logger.info( + "EPD encode embedding cache: L1(VRAM)=%d MiB, L2(host DRAM)=%d MiB", + l1_bytes >> 20, + l2_bytes >> 20, + ) + return TieredEmbeddingCache(l1_bytes, l2_bytes, device=device) + logger.info( + "EPD encode embedding cache: L1(VRAM)=%d MiB (host-DRAM L2 disabled)", + l1_bytes >> 20, + ) + return EmbeddingCache(l1_bytes) + + +def _build_manager_args(server_args, mapping): + """EmbeddingManagerArgs for the encode Mooncake endpoint. + + EPD embedding transfer currently supports TP only. Horizontal scale is via + separate encode servers selected by the gateway; attention DP inside one + encode server is rejected so it cannot be accidentally advertised as TP. + """ + from tokenspeed.runtime.pd.epd.entities import EmbeddingManagerArgs + + if mapping.attn.dp_size != 1: + raise ValueError( + "disaggregation_mode=encode currently supports encode " + f"data_parallel_size == 1, got {mapping.attn.dp_size}" + ) + bootstrap_host = None + if server_args.dist_init_addr: + bootstrap_host = server_args.dist_init_addr.split(":", 1)[0] + return EmbeddingManagerArgs( + bootstrap_port=server_args.disaggregation_bootstrap_port, + tp_size=mapping.attn.tp_size, + bootstrap_host=bootstrap_host, + ) + + +def _maybe_install_encoder_cudagraph(model, server_args) -> bool: + """Install the vision-encoder CUDA-graph wrapper as ``model.image_encoder``, + mirroring the aggregated path's hook in ``execution/model_executor.py``. + + The encode loop never builds a ModelExecutor, so this is where the encode + worker opts into capture/replay of the tower instead of running it eager. Same + gate as the aggregated install: the model exposes the builder, multimodal is + active, the env flag is on, and the attention backend is graph-capturable. The + wrapper IS the model's ``image_encoder`` seam (lazy capture on first encode); + the executor's IMAGE path dispatches through ``model.image_encoder`` and falls + back to eager ``get_image_feature`` (the default) when this returns False. + Returns whether the wrapper was installed. + """ + if not ( + hasattr(model, "make_encoder_cudagraph_wrapper") + and getattr(model, "is_multimodal_active", True) + and envs.TOKENSPEED_MM_ENABLE_ENCODER_CUDA_GRAPH.get() + and server_args.mm_attention_backend != "flashinfer_cudnn" + ): + return False + model.image_encoder = model.make_encoder_cudagraph_wrapper(model.mapping) + logger.info("EPD encode worker: vision-encoder CUDA graph installed") + return True + + +def _build_encode_worker(server_args, port_args, gpu_id, global_rank): + """Assemble the encode worker: model + Mooncake manager + bootstrap server + + executor + scheduler + cache, driven from the real ServerArgs.""" + from tokenspeed.runtime.configs.model_config import ModelConfig + from tokenspeed.runtime.execution.distributed_initializer import ( + DistributedConfig, + DistributedInitializer, + ) + from tokenspeed.runtime.execution.factory import create_model_runner + from tokenspeed.runtime.pd.epd.conn import MooncakeEmbeddingBootstrapServer + from tokenspeed.runtime.pd.epd.embedding_transfer import ( + MooncakeEmbeddingManagerEncode, + ) + from tokenspeed.runtime.pd.epd.encode_executor import ( + DisaggEncodeExecutor, + ) + from tokenspeed.runtime.pd.epd.entities import EmbeddingArgs + + mapping = server_args.mapping + attn_tp_rank = mapping.attn.tp_rank + device = f"cuda:{gpu_id}" + + model_config = ModelConfig( + server_args.model, + trust_remote_code=server_args.trust_remote_code, + revision=server_args.revision, + context_length=server_args.max_model_len, + model_override_args=server_args.hf_overrides, + dtype=server_args.dtype, + quantization=server_args.quantization, + server_args=server_args, + ) + DistributedInitializer.initialize( + DistributedConfig.from_server_args( + server_args=server_args, + port_args=port_args, + gpu_id=gpu_id, + global_rank=global_rank, + hidden_size=model_config.hidden_size, + max_num_tokens=server_args.chunked_prefill_size or 8192, + ) + ) + # The encode worker only needs the vision tower. The model is built + # vision-only (LM construction + LM weight load skipped) via the + # encoder_only gate derived from disaggregation_mode=="encode". The tower + # is used via DisaggEncodeExecutor. No KV/mamba pool is allocated: the encode + # loop never builds a ModelExecutor. + model = create_model_runner(server_args, model_config, None, gpu_id, global_rank)[ + 0 + ].model + # Opt into vision-encoder CUDA-graph capture (mirrors the aggregated path, + # which the encode worker bypasses by never building a ModelExecutor). + _maybe_install_encoder_cudagraph(model, server_args) + + manager_args = _build_manager_args(server_args, mapping) + embedding_args = EmbeddingArgs( + engine_rank=global_rank, + gpu_id=gpu_id, + ib_device=server_args.disaggregation_ib_device, + embedding_data_ptr=0, + embedding_data_len=0, + ) + # The encode worker is the Mooncake data source: it hosts its own bootstrap + # server (prefill workers discover it via the handshake the gateway injects). + # At TP>1 only rank 0 binds the port; every rank still registers its own + # rank_ip/rank_port to the rank-0 bootstrap host, so each prefill rank can + # look up the encode rank it pairs with (contiguous blocks of prefill ranks + # share one encode rank; encode_tp=1 -> all prefill ranks pair encode rank 0). + if attn_tp_rank == 0: + MooncakeEmbeddingBootstrapServer(server_args.disaggregation_bootstrap_port) + manager = MooncakeEmbeddingManagerEncode(manager_args, embedding_args) + executor = DisaggEncodeExecutor(manager, model, device=device) + + l1_bytes = _embedding_cache_bytes(envs.TOKENSPEED_EPD_ENCODE_EMBED_CACHE_MB) + l2_bytes = _embedding_cache_bytes(envs.TOKENSPEED_EPD_ENCODE_EMBED_CACHE_DRAM_MB) + cache = _make_embedding_cache(l1_bytes, l2_bytes, device) + scheduler = EncodeScheduler( + max_tokens_per_batch=server_args.chunked_prefill_size or 8192, + max_items_per_batch=server_args.max_num_seqs, + ) + return EncodeWorker(executor, scheduler, cache), model_config + + +def run_encode_loop(server_args, port_args, pipe_writer, gpu_id, global_rank): + """Run the encode-worker loop until the parent process exits. + + Drains :class:`EncodeRequest`s off the shared scheduler-input ZMQ channel into + EncodeWorker.submit, then runs EncodeWorker.step to encode + ship each pending + item over Mooncake. Synchronous; no KV, no LM forward. + """ + worker, model_config = _build_encode_worker( + server_args, port_args, gpu_id, global_rank + ) + + # TP coordination: the vision tower is TP-sharded (collective ops), so every + # encode rank must run each batch in lockstep. Only rank 0 owns the gateway + # ZMQ; it broadcasts each drained batch to the TP group so all ranks submit + # the same requests and step together. (TP=1 -> rank 0 only, no broadcast.) + attn_tp_rank = server_args.mapping.attn.tp_rank + attn_tp_size = server_args.attn_tp_size or server_args.mapping.attn.tp_size + tp_cpu_group = None + broadcast_pyobj = None + attn_tp_src_rank = server_args.mapping.attn.tp_group[0] + if attn_tp_size > 1: + from tokenspeed.runtime.distributed.process_group_manager import ( + process_group_manager as pg_manager, + ) + from tokenspeed.runtime.utils.common import broadcast_pyobj as _bcast + + broadcast_pyobj = _bcast + tp_cpu_group = pg_manager.get_process_group( + "gloo", server_args.mapping.attn.tp_group + ) + + context = zmq.Context(2) + recv_from_gateway = None + if attn_tp_rank == 0: + recv_from_gateway = get_zmq_socket( + context, zmq.PULL, port_args.scheduler_input_ipc_name, False + ) + + # Unblock the launcher. The encode role has no KV pool, so the token/seq + # fields are nominal (kept for envelope compatibility with the LM ready msg). + pipe_writer.send( + { + "status": "ready", + "max_total_num_tokens": 0, + "max_req_input_len": model_config.context_len, + "max_num_seqs": server_args.max_num_seqs, + "chunked_prefill_size": server_args.chunked_prefill_size, + "max_model_len": model_config.context_len, + } + ) + + while True: + # Rank 0 drains the gateway ZMQ without blocking; other ranks get the + # same batch by broadcast below. + new_reqs = [] + if attn_tp_rank == 0: + while True: + try: + request = recv_from_gateway.recv_pyobj(flags=zmq.NOBLOCK) + except zmq.Again: + break + new_reqs.append(request) + + if attn_tp_size > 1: + # Unconditional every iteration: this is the TP rendezvous that keeps + # all ranks stepping the (collective) vision tower in lockstep. + new_reqs = broadcast_pyobj( + new_reqs, attn_tp_rank, tp_cpu_group, src=attn_tp_src_rank + ) + + for request in new_reqs: + worker.submit(request) + drained = len(new_reqs) > 0 + + # Encode + ship one scheduler batch. step() returns 0 when idle. All ranks + # hold the same pending set, so they make the same step decision together. + did = worker.step() + + # Yield the GIL when there is no fresh ZMQ work, OR when sends are deferred + # on a full ring. The daemon transfer-workers that free ring slots are + # GIL-bound, so spinning here would starve them and wedge the ring; a brief + # yield lets them free slots so the next drain_deferred() can ship. + if ( + not drained and did == 0 and not worker.has_pending() + ) or worker.has_deferred(): + time.sleep(0.0005) diff --git a/python/tokenspeed/runtime/pd/epd/encode_scheduler.py b/python/tokenspeed/runtime/pd/epd/encode_scheduler.py new file mode 100644 index 0000000..9555a7d --- /dev/null +++ b/python/tokenspeed/runtime/pd/epd/encode_scheduler.py @@ -0,0 +1,113 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Encode-stage batching for EPD transfer. + +The encode server runs the vision tower only, orchestrated in Python rather +than the C++ KV scheduler, so all it needs is to batch pending vision items +into a single ViT forward under a token budget -- that is +:class:`EncodeScheduler` below. + +The duplicate-image cache lives in +:mod:`tokenspeed.runtime.cache.embedding_cache`. The two are decoupled: the +encode loop checks the cache on arrival and feeds only *misses* to the +scheduler. +""" + +from __future__ import annotations + +import dataclasses + + +@dataclasses.dataclass(frozen=True) +class PendingEncodeItem: + """One vision item awaiting encode, identified within its request. + + ``cost`` is the item's vision-token (patch) count, used as the batching + budget unit. + """ + + request_id: str + item_index: int + cost: int + + @property + def key(self) -> tuple[str, int]: + return (self.request_id, self.item_index) + + +class EncodeScheduler: + """Deterministic patch-budget batcher for the encode (vision-tower) stage. + + Collects pending vision items (cache misses) and packs them into batches + bounded by a per-batch token budget and a max item count. Ordering is + deterministic across tensor-parallel ranks -- items are sorted by + ``(request_id, item_index)`` -- so every rank forms an identical ViT batch. + This matters because the vision tower can be tensor-parallel; non-identical + batches across ranks would deadlock the NCCL collectives inside it. + + Greedy packing: items are admitted in order until the next one would exceed + ``max_tokens_per_batch`` or ``max_items_per_batch``. A single item whose + cost alone exceeds the token budget is still returned alone, so it always + makes progress. + """ + + def __init__(self, max_tokens_per_batch: int, max_items_per_batch: int): + if max_tokens_per_batch <= 0: + raise ValueError( + f"max_tokens_per_batch must be > 0, got {max_tokens_per_batch}" + ) + if max_items_per_batch <= 0: + raise ValueError( + f"max_items_per_batch must be > 0, got {max_items_per_batch}" + ) + self.max_tokens_per_batch = max_tokens_per_batch + self.max_items_per_batch = max_items_per_batch + self._pending: dict[tuple[str, int], PendingEncodeItem] = {} + + def add(self, item: PendingEncodeItem) -> None: + # Idempotent on (request_id, item_index): a re-added item overwrites. + self._pending[item.key] = item + + def pending_size(self) -> int: + return len(self._pending) + + def _ordered_pending(self) -> list[PendingEncodeItem]: + # Sort by (request_id, item_index) for cross-rank determinism. + return [self._pending[k] for k in sorted(self._pending.keys())] + + def next_batch(self) -> list[PendingEncodeItem]: + """Pop and return the next deterministic batch of items to encode. + + Empty when nothing is pending. Removes the returned items from the + pending set. + """ + batch: list[PendingEncodeItem] = [] + used = 0 + for it in self._ordered_pending(): + if batch and ( + used + it.cost > self.max_tokens_per_batch + or len(batch) >= self.max_items_per_batch + ): + break + batch.append(it) + used += it.cost + del self._pending[it.key] + return batch diff --git a/python/tokenspeed/runtime/pd/epd/encode_worker.py b/python/tokenspeed/runtime/pd/epd/encode_worker.py new file mode 100644 index 0000000..b800bfc --- /dev/null +++ b/python/tokenspeed/runtime/pd/epd/encode_worker.py @@ -0,0 +1,191 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Encode-worker control loop for EPD (Python orchestration). + +This is the body the engine's encode event loop drives: it sits between request +arrival and the vision tower. On ``submit`` it registers the request's transfer +peer and, per item, either resolves the embedding from the cache (skip the +tower, still transfer) or queues it on the scheduler. Each ``step`` pulls one +deterministic batch off the scheduler, runs the tower + ships it via the +executor, and populates the cache. + +The model load, mooncake manager construction, request transport and the +event-loop wiring are supplied by the engine integration; this class only +orchestrates them, so it is unit-testable with fakes. +""" + +from __future__ import annotations + +import dataclasses + +from tokenspeed.runtime.cache.embedding_cache import ( + EmbeddingCache, + TieredEmbeddingCache, +) +from tokenspeed.runtime.multimodal.embedder import _item_token_count +from tokenspeed.runtime.multimodal.inputs import MultimodalDataItem +from tokenspeed.runtime.multimodal.shm_transport import ShmTensorHandle +from tokenspeed.runtime.pd.epd.encode_scheduler import ( + EncodeScheduler, + PendingEncodeItem, +) +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +@dataclasses.dataclass(frozen=True) +class EncodeRequest: + """One encode request: a transfer peer plus its vision items. + + ``bootstrap_host``/``port``/``room`` identify the prefill peer this + request's embeddings are shipped to (assigned upstream, per request). + """ + + request_id: str + bootstrap_host: str + bootstrap_port: int + bootstrap_room: int + items: list[MultimodalDataItem] + + +def _nbytes(tensor) -> int: + return tensor.numel() * tensor.element_size() + + +class EncodeWorker: + """Orchestrates cache + scheduler + executor for the encode role. + + Injected with the executor (real ``DisaggEncodeExecutor`` or a fake), an + ``EncodeScheduler`` and an embedding cache (single-tier ``EmbeddingCache`` or + the two-tier ``TieredEmbeddingCache``; only ``get``/``put`` are used) so the + control flow is testable without a GPU or transport. + """ + + def __init__( + self, + executor, + scheduler: EncodeScheduler, + cache: EmbeddingCache | TieredEmbeddingCache, + ): + self.executor = executor + self.scheduler = scheduler + self.cache = cache + # (request_id, item_index) -> item awaiting the tower + self._pending: dict = {} + + def submit(self, request: EncodeRequest) -> None: + self.executor.register( + request.request_id, + request.bootstrap_host, + request.bootstrap_port, + request.bootstrap_room, + ) + for idx, item in enumerate(request.items): + cached = self.cache.get(item.hash) + if isinstance(item.feature, ShmTensorHandle): + # EPD pixel-SHM: the servicer published pixels to POSIX SHM and + # the ZMQ hop carried only this handle (hash/pad_value were set on + # the real tensor before publish). consume() unlinks, so segments + # never outlive the item: materialize on a miss, and on a hit still + # consume-and-drop to unlink the unused segment. + handle, item.feature = item.feature, None + handle.attach() + if cached is None: + item.feature = handle.consume() + else: + handle.consume() + if cached is not None: + # Cache hit: tower skipped, but the embedding still must reach + # the prefill peer, so ship it directly. Entries are + # (main, deepstack) pairs and BOTH halves must be restored, else + # the prefill publishes a never-written deepstack buffer. Tolerate + # a bare tensor for legacy/test-seeded entries. + if isinstance(cached, tuple): + item.encoded, item.encoded_deepstack = cached + else: + item.encoded = cached + self.executor.send_item(request.request_id, item) + else: + self.scheduler.add( + PendingEncodeItem( + request_id=request.request_id, + item_index=idx, + cost=_item_token_count(item), + ) + ) + self._pending[(request.request_id, idx)] = item + + def step(self) -> int: + """Run one scheduler batch through the tower + transfer. Returns the + number of items encoded (0 when nothing is pending).""" + self.executor.reap_concluded_senders({rid for (rid, _idx) in self._pending}) + # Retry sends that couldn't lease a ring slot last tick (non-blocking). + self.executor.drain_deferred() + # Backpressure: if sends are STILL deferred after the drain, the bounce + # ring is saturated (all slots hold in-flight transfers). Pulling more ViT + # now would only pile fresh embeddings into _deferred_sends -- each pins a + # GPU embedding tensor with no slot to ship it, growing an unbounded + # backlog into an OOM. Skip this tick; the loop yields the GIL (encode_loop + # sees has_deferred) so the transfer daemons free slots, then we resume. + if self.executor.has_deferred(): + return 0 + batch = self.scheduler.next_batch() + if not batch: + return 0 + request_items = [(p.request_id, self._pending[p.key]) for p in batch] + try: + self.executor.execute(request_items) + except Exception as e: + # A tower-step contract violation (ViT output not matching the items' + # post-merge token count, or the forward itself) must fail only the + # rooms in THIS batch, not propagate out into the engine's SIGUSR1 + # handler, which would kill the worker and drop every other request's + # in-flight image. These raises fire before any send is issued, so + # concluding the batch Failed never poisons an already-shipped room. + # Per-item STAGING errors are handled finer-grained inside + # _stage_and_send -> _fail_staged_room. + n_failed = self.executor.fail_rooms((rid for rid, _ in request_items), e) + for p in batch: + self._pending.pop(p.key, None) + logger.error( + "encode batch failed (%d rooms concluded Failed): %s", n_failed, e + ) + return 0 + for p in batch: + item = self._pending.pop(p.key) + if item.encoded is not None: + # Cache the (main, deepstack) PAIR: caching only the main half + # would make every later hit ship a deepstack-less transfer, + # publishing uninitialized rows on the prefill. + deep = item.encoded_deepstack + nbytes = _nbytes(item.encoded) + ( + _nbytes(deep) if deep is not None else 0 + ) + self.cache.put(item.hash, (item.encoded, deep), nbytes) + return len(batch) + + def has_pending(self) -> bool: + return self.scheduler.pending_size() > 0 + + def has_deferred(self) -> bool: + """True while sends are queued waiting for a free ring slot (executor).""" + return self.executor.has_deferred() diff --git a/python/tokenspeed/runtime/pd/epd/entities.py b/python/tokenspeed/runtime/pd/epd/entities.py new file mode 100644 index 0000000..e6e17d3 --- /dev/null +++ b/python/tokenspeed/runtime/pd/epd/entities.py @@ -0,0 +1,249 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import dataclasses +import struct +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import torch + + +# Sentinel room (on-wire ascii) for the one-shot registration frame; per-request +# pre-alloc frames carry an integer room. +REGISTER_ROOM_SENTINEL = "None" + + +def _pack_ptr(ptr: int) -> bytes: + """Encode a device virtual address as 8 little-endian bytes (uint64).""" + return struct.pack("Q", ptr) + + +def _unpack_ptr(frame: bytes) -> int: + return struct.unpack("Q", frame)[0] + + +def _b(value: object) -> bytes: + return str(value).encode("ascii") + + +def _bool_frame(value: bool) -> bytes: + return b"1" if value else b"0" + + +def _parse_bool(frame: bytes) -> bool: + return frame == b"1" + + +@dataclasses.dataclass +class EmbeddingArgs: + """Per-rank buffer/engine info for an embedding endpoint. + + On the encode (sender) side the buffers hold the vision tower's output; on + the prefill (receiver) side they are the pre-registered receive buffers. A + ``deepstack`` buffer is present only for models that emit deepstack + embeddings (e.g. Qwen3.5); ``0`` means absent. + """ + + engine_rank: int + gpu_id: int + ib_device: str | None + embedding_data_ptr: int + embedding_data_len: int + deepstack_data_ptr: int = 0 + deepstack_data_len: int = 0 + + +class EmbeddingTransferError(Exception): + def __init__( + self, bootstrap_room: int, failure_reason: str, remote_endpoint: str = None + ): + super().__init__(failure_reason) + self.bootstrap_room = bootstrap_room + self.failure_reason = failure_reason + self.remote_endpoint = remote_endpoint + + def __str__(self): + if self.remote_endpoint: + return ( + "EmbeddingTransferError(" + f"bootstrap_room={self.bootstrap_room}, " + f"remote_endpoint={self.remote_endpoint}): {self.failure_reason}" + ) + else: + return ( + f"EmbeddingTransferError(bootstrap_room={self.bootstrap_room}): " + f"{self.failure_reason}" + ) + + +@dataclasses.dataclass +class EmbeddingChunk: + """One item's embedding queued for transfer (encode-side, in-process only). + + Not serialized to ZMQ: it lives on the encode manager's transfer queue. The + pointers reference the contiguous ``item.encoded`` / ``item.encoded_deepstack`` + tensors produced by the vision tower. + """ + + room: int + src_embedding_ptr: int + n_tokens: int + hidden: int + dtype: str + nbytes: int + src_deepstack_ptr: int = 0 + deepstack_width: int = 0 # hidden * num_deepstack; 0 == no deepstack + deepstack_nbytes: int = 0 + # CUDA event recorded on the encode loop's stream after the device copy that + # filled this chunk's ring slot. The transfer worker waits it before its + # one-sided RDMA read so the read never races the copy (ViT->send corruption + # hazard). None on CPU/no-CUDA. + copy_event: "torch.cuda.Event | None" = None + + +@dataclasses.dataclass +class EmbeddingTransferInfo: + """Receiver(prefill)->sender(encode) per-request pre-allocation frame. + + Tells the encode side where (``dst_embedding_ptr`` already includes this + request's row offset) and how big a transfer to issue. The RDMA write is + unchecked, so the encode asserts ``n_tokens``/``hidden``/``dtype`` against + the tensor it sends, else a mismatch silently writes a truncated/oversized + region. + + Row sharding: ``row_start`` + ``n_tokens`` (the SHARD's row count) select + the rows the encode writes at the shard-offset ``dst_embedding_ptr``; + ``span`` is the image's FULL row count, the cross-side token-count tripwire + regardless of shard geometry. Identity frames set ``row_start == 0`` and + ``n_tokens == span``. The encode validates shard geometry per frame: a + full-span write at a shard-offset pointer would corrupt neighboring image + rows in the same buffer. + """ + + room: int + endpoint: str + dst_port: int + mooncake_session_id: str + dst_embedding_ptr: int + dst_deepstack_ptr: int + n_tokens: int + hidden: int + dtype: str + has_deepstack: bool + required_dst_info_num: int + row_start: int = 0 + span: int = 0 + + def to_zmq(self) -> list[bytes]: + return [ + _b(self.room), + _b(self.endpoint), + _b(self.dst_port), + _b(self.mooncake_session_id), + _pack_ptr(self.dst_embedding_ptr), + _pack_ptr(self.dst_deepstack_ptr), + _b(self.n_tokens), + _b(self.hidden), + _b(self.dtype), + _bool_frame(self.has_deepstack), + _b(self.required_dst_info_num), + _b(self.row_start), + _b(self.span), + ] + + @classmethod + def from_zmq(cls, msg: list[bytes]) -> "EmbeddingTransferInfo": + # Length-guarded so a truncated frame is logged and dropped by the + # bootstrap listener instead of mis-parsing. + if len(msg) < 13: + raise ValueError( + f"malformed EmbeddingTransferInfo frame: {len(msg)} parts < 13" + ) + 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_embedding_ptr=_unpack_ptr(msg[4]), + dst_deepstack_ptr=_unpack_ptr(msg[5]), + n_tokens=int(msg[6].decode("ascii")), + hidden=int(msg[7].decode("ascii")), + dtype=msg[8].decode("ascii"), + has_deepstack=_parse_bool(msg[9]), + required_dst_info_num=int(msg[10].decode("ascii")), + row_start=int(msg[11].decode("ascii")), + span=int(msg[12].decode("ascii")), + ) + + +@dataclasses.dataclass +class EmbeddingArgsRegisterInfo: + """Receiver(prefill)->sender(encode) one-shot endpoint registration. + + Sent once per receiver rank with room == ``REGISTER_ROOM_SENTINEL`` so the + encode side records where to PUSH completion status. The buffer pointers + are the receiver's base receive buffers (per-item row offsets are carried + later, in :class:`EmbeddingTransferInfo`). + """ + + room: str + endpoint: str + dst_port: int + mooncake_session_id: str + dst_embedding_ptr: int + dst_deepstack_ptr: int + + def to_zmq(self) -> list[bytes]: + return [ + _b(self.room), + _b(self.endpoint), + _b(self.dst_port), + _b(self.mooncake_session_id), + _pack_ptr(self.dst_embedding_ptr), + _pack_ptr(self.dst_deepstack_ptr), + ] + + @classmethod + def from_zmq(cls, msg: list[bytes]) -> "EmbeddingArgsRegisterInfo": + return cls( + room=msg[0].decode("ascii"), + endpoint=msg[1].decode("ascii"), + dst_port=int(msg[2].decode("ascii")), + mooncake_session_id=msg[3].decode("ascii"), + dst_embedding_ptr=_unpack_ptr(msg[4]), + dst_deepstack_ptr=_unpack_ptr(msg[5]), + ) + + +@dataclasses.dataclass +class EmbeddingManagerArgs: + """Connection config for the embedding Mooncake endpoint. + + Embedding transfer has a single DP group today. ``tp_size`` is the number + of encode/prefill endpoint ranks that participate in that group; the shared + bootstrap layer still receives it as ``world_size`` with ``dp_size=1``. + """ + + bootstrap_port: int + tp_size: int + bootstrap_host: str | None = None diff --git a/python/tokenspeed/runtime/pd/epd/prefill_receiver.py b/python/tokenspeed/runtime/pd/epd/prefill_receiver.py new file mode 100644 index 0000000..4b41924 --- /dev/null +++ b/python/tokenspeed/runtime/pd/epd/prefill_receiver.py @@ -0,0 +1,1092 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""EPD prefill-side receive glue: fill each multimodal item's ``encoded`` +tensor from the Mooncake transfer instead of running the vision tower. + +Inverse of ``encode_executor.assign_encoded_embeddings``: the encode worker +row-splits the tower output per item and column-splits the deepstack half; here +the prefill allocates a receive buffer sized to each item's post-merge token +count (and deepstack width), registers it with that item's encode worker via +:class:`MooncakeEmbeddingReceiver`, waits for the push, and assigns the result +onto ``item.encoded`` / ``item.encoded_deepstack`` -- the ``skip-ViT`` form the +prefill ``VisionEmbedder`` consumes (it never re-encodes an item whose +``encoded`` is already set). A wrong token count or deepstack width mis-sizes the +RDMA target (silent corruption), so buffer-sizing + assignment is isolated and +CPU-unit-testable with a fake receiver. + +Poll-driven state machine (:class:`EmbeddingReceiveJob`): ``start()`` does only +the non-blocking Phase-1 allocation + receiver construction (it must NEVER block +the prefill event loop); ``poll()`` advances every receiver a little each cycle +and publishes ``item.encoded`` only once EVERY handshaked item has landed. The +scheduler holds the job and does not admit the request to a prefill forward until +``poll()`` returns ``DONE``. + +``receive_encoded_embeddings`` is a thin BLOCKING wrapper (start + spin poll) for +synchronous callers and the CPU tests. +""" + +from __future__ import annotations + +import logging +import time +from collections import deque +from collections.abc import Callable, Mapping, Sequence +from typing import Any + +import torch +import torch.distributed as dist + +logger = logging.getLogger(__name__) + +from tokenspeed.runtime.multimodal.embedder import _item_token_count +from tokenspeed.runtime.multimodal.inputs import MultimodalDataItem +from tokenspeed.runtime.pd.base.status import TransferPoll +from tokenspeed.runtime.pd.epd.embedding_transfer import ( + MooncakeEmbeddingReceiver, +) +from tokenspeed.runtime.utils.env import envs + +# (manager, bootstrap_addr, bootstrap_room) -> receiver. Defaults to the real +# MooncakeEmbeddingReceiver; overridden in tests with a fake. +ReceiverFactory = Callable[[Any, str, int], Any] + +# poll() return values (the per-job lifecycle status, distinct from the +# per-receiver Poll status it aggregates). +PENDING = "pending" +DONE = "done" +FAILED = "failed" + +# Deregistering a receive MR the instant its Success notif arrives races the +# NIC's DMA placement tail (Success is the transport ACK, not proof the last +# bytes cleared HCA->PCIe), tripping `local access violation work queue error`. +# So deregistration is DEFERRED: entries hold the tensor ref (the allocator +# cannot reuse a still-registered address -- the no-double-register invariant) +# and are swept after a grace period, on the scheduler loop (no thread, no lock). +_DEREG_DELAY_S = 0.5 +_RECV_POOL_QUARANTINE_S = 10.0 +# (due_monotonic, engine, [(tensor, ptr), ...]) in due order (delay is constant). +_pending_dereg: deque = deque() + + +def _record_current_stream_event(tensor: torch.Tensor) -> torch.cuda.Event | None: + if not tensor.is_cuda: + return None + event = torch.cuda.Event() + torch.cuda.current_stream(tensor.device).record_event(event) + return event + + +def _lazy_deregister(engine: Any, tensors: list[tuple[torch.Tensor, int]]) -> None: + _pending_dereg.append((time.monotonic() + _DEREG_DELAY_S, engine, tensors)) + _sweep_deregister() + + +def _sweep_deregister() -> None: + now = time.monotonic() + while _pending_dereg and _pending_dereg[0][0] <= now: + _, engine, tensors = _pending_dereg.popleft() + for _tensor, ptr in tensors: + try: + engine.deregister(ptr) + except Exception: # noqa: BLE001 -- best-effort; worst case the MR leaks + pass + + +class _RecvBufferPool: + """Pre-registered, lifetime-stable receive slots for E->P embedding lands. + + With per-request buffers, deregistering then re-registering a recycled + allocator address mints a NEW rkey for the same range while the encode side's + Mooncake segment cache still resolves it to the OLD rkey -- every reuse risks + a `local access violation work queue error` that kills the QP. The pool + registers ONE region for the engine's lifetime, so the sender's cached + mapping can never go stale; requests lease slots, the publish path clones the + landed rows out, and CUDA-backed slots return only after that clone completes. + + Failure path: a FAILED job may still have an in-flight remote write targeting + its slot, which under a single lifetime MR would land SILENTLY in the next + tenant's data, so failed slots sit in quarantine until the transfer layer's + timeouts have expired. + + Single-threaded by design: all callers run on the scheduler loop (no locks). + """ + + def __init__(self, engine: Any, device: Any, slot_bytes: int, n_slots: int): + self.engine = engine + self.slot_bytes = slot_bytes + self.buf = torch.empty(n_slots * slot_bytes, dtype=torch.uint8, device=device) + engine.register(self.buf.data_ptr(), self.buf.numel()) + self._free = list(range(n_slots)) + self._quarantine: deque = deque() # (release_due_monotonic, slot) + self._pending_release: deque = deque() # (cuda_event, slot) + + def _sweep_pending_release(self) -> None: + kept = deque() + while self._pending_release: + event, slot = self._pending_release.popleft() + try: + ready = event.query() + except Exception: # noqa: BLE001 -- avoid permanently losing a slot + logger.warning( + "EPD recv pool: CUDA event query failed; releasing slot anyway", + exc_info=True, + ) + ready = True + if ready: + self._free.append(slot) + else: + kept.append((event, slot)) + self._pending_release = kept + + def _sweep_quarantine(self) -> None: + now = time.monotonic() + while self._quarantine and self._quarantine[0][0] <= now: + self._free.append(self._quarantine.popleft()[1]) + + def sweep(self) -> None: + self._sweep_pending_release() + self._sweep_quarantine() + + def lease(self, nbytes: int) -> int | None: + self.sweep() + if nbytes > self.slot_bytes or not self._free: + return None + return self._free.pop() + + def view(self, slot: int, nbytes: int, dtype: torch.dtype, shape) -> torch.Tensor: + off = slot * self.slot_bytes + return self.buf[off : off + nbytes].view(dtype).reshape(shape) + + def release(self, slot: int) -> None: + self._free.append(slot) + + def release_after_copy(self, slot: int, copied_tensor: torch.Tensor) -> None: + event = _record_current_stream_event(copied_tensor) + if event is None: + self.release(slot) + else: + self._pending_release.append((event, slot)) + + def quarantine(self, slot: int, delay_s: float) -> None: + self._quarantine.append((time.monotonic() + delay_s, slot)) + + +# (id(engine), str(device)) -> _RecvBufferPool | False (False = disabled). +_POOLS: dict = {} + + +def _get_pool(engine: Any, device: Any) -> _RecvBufferPool | None: + key = (id(engine), str(device)) + pool = _POOLS.get(key) + if pool is None: + n_slots = envs.TOKENSPEED_EPD_RECV_POOL_SLOTS.get() + slot_mb = envs.TOKENSPEED_EPD_RECV_POOL_SLOT_MB.get() + if n_slots <= 0 or slot_mb <= 0: + pool = False + else: + pool = _RecvBufferPool(engine, device, slot_mb << 20, n_slots) + logger.info( + "EPD recv pool up: %d slots x %d MB (lifetime MR)", n_slots, slot_mb + ) + _POOLS[key] = pool + return pool or None + + +def shard_rows(span: int, shard_rank: int, shard_size: int) -> tuple[int, int]: + """Balanced contiguous row shard of ``span`` rows across ``shard_size`` + ranks: returns this rank's ``(row_start, row_count)``. The first + ``span % shard_size`` ranks get one extra row; shards tile ``[0, span)`` + disjointly and in rank order, which BOTH sides of the transfer (the + receiver's pre_alloc and the post-receive reassembly) must derive from this + one function so their geometry can never diverge. ``row_count`` may be 0 + when ``span < shard_size`` (tiny images).""" + base, rem = divmod(span, shard_size) + start = shard_rank * base + min(shard_rank, rem) + count = base + (1 if shard_rank < rem else 0) + return start, count + + +class _ItemReceive: + """Per-item receive bookkeeping for one :class:`EmbeddingReceiveJob`. + + Holds the single ``[n_tokens, hidden]`` receive buffer (+ optional deepstack + column buffer) and one receiver per owned image, each targeting a contiguous + row sub-range of that buffer. The buffers are filled in place by the encode + side's RDMA writes; this object publishes them onto the item only once every + one of its receivers reaches ``Success`` (handled by the owning job). + """ + + __slots__ = ( + "item", + "recv_main", + "recv_deepstack", + "receivers", + "pre_alloced", + "main_slice_ptrs", + "deepstack_slice_ptrs", + "spans", + "row_starts", + "row_counts", + "sharded", + "n_tokens", + "hidden", + "pool", + "pool_slot", + ) + + def __init__( + self, + item: MultimodalDataItem, + recv_main: torch.Tensor, + recv_deepstack: torch.Tensor | None, + receivers: list[Any], + main_slice_ptrs: list[int], + deepstack_slice_ptrs: list[int], + spans: list[int], + row_starts: list[int], + row_counts: list[int], + sharded: list[bool], + n_tokens: int, + hidden: int, + pool: _RecvBufferPool | None = None, + pool_slot: int | None = None, + ): + self.item = item + self.recv_main = recv_main + self.recv_deepstack = recv_deepstack + self.receivers = receivers + self.pool = pool + self.pool_slot = pool_slot + # Lazy-pre_alloc latch, one per receiver (index-aligned): each receiver's + # pre_alloc must be issued exactly once, AFTER it reports Bootstrapped, and + # never again (a second pre_alloc would double-register the transfer). + self.pre_alloced = [False] * len(receivers) + self.main_slice_ptrs = main_slice_ptrs + self.deepstack_slice_ptrs = deepstack_slice_ptrs + self.spans = spans + # Per-image shard geometry (index-aligned with receivers/spans): this + # rank's row sub-range WITHIN the image, and whether the image was + # sharded at all (identity images need no reassembly broadcast). + self.row_starts = row_starts + self.row_counts = row_counts + self.sharded = sharded + self.n_tokens = n_tokens + self.hidden = hidden + + +class EmbeddingReceiveJob: + """Poll-driven, non-blocking EPD embedding receive for ONE request. + + Usage (driven by the scheduler/event loop): + + job = start_embedding_receive(items, manager, ...) + # ... each cycle, until DONE/FAILED: + status = job.poll() # PENDING | DONE | FAILED + + ``start()`` (the constructor) runs Phase-1 only: it sizes + allocates each + handshaked item's receive buffer, registers it with the Mooncake engine, and + constructs one :class:`MooncakeEmbeddingReceiver` on the item's room (the + receiver's own ``__init__`` performs the bootstrap handshake). It does NOT + wait for any transfer and does NOT call ``pre_alloc`` -- that is deferred to + ``poll()``, which issues each ``pre_alloc`` lazily once its receiver reports + Bootstrapped and then waits (across cycles, never blocking) for Success. + + HARD CONSTRAINT: one room/receiver PER ITEM. A single request's + items may be served by DIFFERENT encode workers, so collapsing to one + receiver per request is not possible. ``poll()`` returns DONE only when every + item's receiver is Success. + + Buffer lifetime: by default the receive target is a leased slot from the + lifetime-registered :class:`_RecvBufferPool`; on DONE the landed rows are + cloned onto ``item.encoded`` and the slot is reused after that copy completes + (no MR churn, see the pool docstring). Deepstack models, oversized items, pool + exhaustion, or ``TOKENSPEED_EPD_RECV_POOL_SLOTS=0`` fall back to a per-request + buffer registered on start and lazily deregistered after publish. Pooled + slots stay leased until the CUDA copy into ``item.encoded`` has completed, + so a following request cannot overwrite rows still being cloned. The GPU cost + per request is roughly ``n_tokens * hidden * dtype.itemsize`` (plus ``* (1 + + num_deepstack)`` with deepstack); the caller should cap the number of + in-flight jobs accordingly. + + Idempotent re-`start`/poll: items whose ``item.encoded`` is already set + (chunked prefill re-runs the receive per Path-4 forward on the same item) are + SKIPPED at start time -- no receiver is constructed for them. + """ + + def __init__( + self, + items: Sequence[MultimodalDataItem], + manager: Any, + *, + hidden: int, + num_deepstack: int, + dtype: torch.dtype, + device: torch.device | str, + receiver_factory: ReceiverFactory | None = None, + shard_rank: int = 0, + shard_size: int = 1, + ): + self.manager = manager + self.hidden = hidden + self.num_deepstack = num_deepstack + # Row sharding across the attn-TP group: with shard_size > 1 this rank + # registers only its shard_rows() sub-range of each image; after the + # rank-agreed DONE the caller MUST run reassemble() (rank-lockstep) to + # rebuild the full rows before any forward consumes item.encoded. + # shard_size <= 1 is a plain full copy (no reassemble). + self._shard_rank = shard_rank + self._shard_size = shard_size + # torch.dtype is used for buffer allocation (no dtype lost in a string + # round-trip); the string is only the encode-side wire contract in pre_alloc. + self.dtype = dtype + self.dtype_str = str(dtype) + self._factory: ReceiverFactory = receiver_factory or MooncakeEmbeddingReceiver + # Terminal latch: once DONE/FAILED, poll() is a cheap no-op returning it + # (the receivers/buffers have been torn down). + self._status: str = PENDING + + # Phase 1 (non-blocking): for each item carrying an encode handshake + # (``item.encode_handshake``), allocate ONE [n_tokens, hidden] buffer + # (+ deepstack columns), register it with the Mooncake engine, and build a + # single receiver on the item's room. The handshake lives ON the item: the + # gateway mints one room per item and the encode worker row-splits the + # concatenated-subgrid embedding per item -- so one item == one room == one + # embedding. + self._items: list[_ItemReceive] = [] + for item in items: + handshake = getattr(item, "encode_handshake", None) + if handshake is None: + # No EPD-routed embedding on this item; leave it for the tower. + continue + + # Chunked prefill re-runs the receive per forward on the SAME item; + # ``encoded`` is set only after Success, so ``encoded is not None`` + # means fully received -- skip it (re-bootstrapping a Success room + # would never re-report Bootstrapped and would stick). Mirrors + # ``VisionEmbedder``, which never re-encodes an encoded item. + if item.encoded is not None: + continue + + self._items.append(self._start_item(item, handshake, device)) + + # Nothing to receive (text-only / all-already-encoded / no EPD item): + # the job is immediately DONE so the scheduler admits the request at once. + if not self._items: + self._status = DONE + + def _start_item( + self, + item: MultimodalDataItem, + handshake: Mapping[str, Any], + device: torch.device | str, + ) -> _ItemReceive: + """Allocate + register one item's receive buffer and construct its single + receiver on the item's room, recording the (possibly sharded) destination + row range so ``poll()`` can lazily issue ``pre_alloc`` once the receiver + bootstraps. Does NOT block on any transfer (no ``_wait``, no ``pre_alloc`` + here). + + One room per item: the encode worker concatenates the item's subgrid + tokens into a single ``[n_tokens, hidden]`` embedding and row-splits it per + prefill TP rank, so the receive geometry spans the item's FULL token count + (sum of its offset spans), not per offset. + """ + elt = torch.empty(0, dtype=self.dtype).element_size() + + # The item's image spans all its concatenated subgrid tokens. + span = _item_token_count(item) + addr = f"{handshake['bootstrap_host']}:{handshake['bootstrap_port']}" + # The receiver __init__ performs the bootstrap handshake (HTTP /route fetch + # + endpoint registration); on return its poll() is already Bootstrapped or + # Failed. We do NOT pre_alloc here -- poll() does it once Bootstrapped. + receiver = self._factory(self.manager, addr, int(handshake["bootstrap_room"])) + + # Row sharding across the attn-TP group: this rank registers only its + # shard_rows() sub-range, or the full image when sharding is off. + is_sharded = self._shard_size > 1 + if is_sharded: + row_start, row_count = shard_rows(span, self._shard_rank, self._shard_size) + else: + row_start, row_count = 0, span + # Length-1 per-image lists so poll()/_packed_to_full()/reassemble() iterate + # unchanged (one item == one image under per-item rooms). + receivers: list[Any] = [receiver] + spans: list[int] = [span] + row_starts: list[int] = [row_start] + row_counts: list[int] = [row_count] + sharded: list[bool] = [is_sharded] + + # The RECEIVE buffer holds only THIS rank's shard rows, packed contiguously + # (image i at rows [packed_cursor, packed_cursor + row_count)); the FULL + # embedding is rebuilt later in item.encoded (publish scatter + reassemble), + # which needs no MR. Registering just the shard rows -- not the full image + # -- keeps a big multi-image item from overflowing a recv pool slot onto + # the GIL-held per-request register fallback. With nothing sharded, + # packed_tokens == full_tokens. + full_tokens = _item_token_count(item) + packed_tokens = sum(row_counts) + nbytes = packed_tokens * self.hidden * elt + pool = None + pool_slot = None + recv_main = None + if self.num_deepstack == 0: + # Preferred path: lease a lifetime-registered pool slot (see + # _RecvBufferPool for why per-request register/deregister kills QPs). + # Deepstack models keep the legacy path (a second buffer per item, + # out of the pool's scope). + pool = _get_pool(self.manager.engine, device) + if pool is not None: + pool_slot = pool.lease(nbytes) + if pool_slot is not None: + recv_main = pool.view( + pool_slot, nbytes, self.dtype, (packed_tokens, self.hidden) + ) + else: + logger.info( + "EPD recv pool: no slot for %d B (free=%d); falling back " + "to per-request registration", + nbytes, + len(pool._free), + ) + pool = None + if recv_main is None: + recv_main = torch.empty( + (packed_tokens, self.hidden), dtype=self.dtype, device=device + ) + # Legacy fallback (pool disabled/exhausted, oversized item, or + # deepstack present): register per request; the deregister is DEFERRED + # on the publish/fail paths (lazy queue) to soften -- not eliminate -- + # the stale-rkey window. A 0-row packed buffer has nothing to register. + if packed_tokens > 0: + self.manager.engine.register( + recv_main.data_ptr(), + recv_main.numel() * recv_main.element_size(), + ) + recv_deepstack: torch.Tensor | None = None + if self.num_deepstack > 0: + recv_deepstack = torch.empty( + (packed_tokens, self.hidden * self.num_deepstack), + dtype=self.dtype, + device=device, + ) + if packed_tokens > 0: + self.manager.engine.register( + recv_deepstack.data_ptr(), + recv_deepstack.numel() * recv_deepstack.element_size(), + ) + + # Destination pointers into the PACKED buffer: the encode side writes its + # n_tokens(=row_count) source rows CONTIGUOUSLY at this pointer (it does + # not re-apply row_start -- that only selects which SOURCE rows to read). + # A 0-row shard still records its (empty) pointer; its pre_alloc is the + # registration heartbeat the encode fanout gate counts. + main_slice_ptrs: list[int] = [] + deepstack_slice_ptrs: list[int] = [] + packed_cursor = 0 + for row_count in row_counts: + main_slice_ptrs.append( + recv_main[packed_cursor : packed_cursor + row_count].data_ptr() + ) + if recv_deepstack is not None: + deepstack_slice_ptrs.append( + recv_deepstack[packed_cursor : packed_cursor + row_count].data_ptr() + ) + else: + deepstack_slice_ptrs.append(0) + packed_cursor += row_count + + return _ItemReceive( + item=item, + recv_main=recv_main, + recv_deepstack=recv_deepstack, + receivers=receivers, + main_slice_ptrs=main_slice_ptrs, + deepstack_slice_ptrs=deepstack_slice_ptrs, + spans=spans, + row_starts=row_starts, + row_counts=row_counts, + sharded=sharded, + n_tokens=full_tokens, + hidden=self.hidden, + pool=pool, + pool_slot=pool_slot, + ) + + def poll(self) -> str: + """Advance the receive state machine one cheap step. + + For every still-pending receiver: if it just reached Bootstrapped, issue + its (single) ``pre_alloc`` so the encode side learns where to write; if it + reached Failed -> the whole job is FAILED; otherwise leave it. When EVERY + receiver of EVERY item is Success, deregister the buffers, publish + ``item.encoded`` / ``item.encoded_deepstack``, and return DONE. + + Returns ``PENDING`` | ``DONE`` | ``FAILED``. Idempotent after a terminal + result (the buffers/receivers are gone, so it just returns the latch). + """ + _sweep_deregister() + if self._status is not PENDING: + return self._status + + all_success = True + for it in self._items: + for idx, receiver in enumerate(it.receivers): + status = receiver.poll() + if status == TransferPoll.Failed: + self._fail() + return FAILED + if not it.pre_alloced[idx] and status in ( + TransferPoll.Bootstrapped, + TransferPoll.Transferring, + TransferPoll.Success, + ): + # Bootstrapped (or already further along on a fast/in-process + # transport): issue the one-shot pre_alloc and latch it -- a + # second would double-register the transfer on the encode side. + # In shard mode n_tokens carries this rank's row COUNT and the + # dst pointers are already offset to the shard's first row; a + # row_count of 0 is still sent (it doubles as the encode-side + # fanout-gate heartbeat). + receiver.pre_alloc( + dst_embedding_ptr=it.main_slice_ptrs[idx], + n_tokens=it.row_counts[idx], + hidden=it.hidden, + dtype=self.dtype_str, + dst_deepstack_ptr=it.deepstack_slice_ptrs[idx], + has_deepstack=self.num_deepstack > 0, + row_start=it.row_starts[idx], + span=it.spans[idx], + ) + it.pre_alloced[idx] = True + # Re-poll once after pre_alloc: an in-process/fake transport + # may flip straight to Success on pre_alloc. + status = receiver.poll() + if status == TransferPoll.Failed: + self._fail() + return FAILED + if status != TransferPoll.Success: + all_success = False + + if not all_success: + return PENDING + + # Every image of every item has landed; publish and reclaim. Pooled + # path: clone the landed rows OUT of the slot (item.encoded must outlive + # the lease) and release the slot only after that copy completes. Legacy + # path: hand the buffer itself to item.encoded and queue the registration + # for DEFERRED drop (the lazy entry holds the tensor ref, so the allocator + # cannot recycle a still-registered address). + for it in self._items: + any_sharded = any(it.sharded) + if it.pool is not None: + # Pooled path is deepstack-free. Copy the rows OUT of the slot + # before releasing it. When sharded the buffer is packed (this + # rank's rows only), so scatter into a full-layout tensor at each + # image's absolute offset (reassemble fills the rest); when not + # sharded the buffer is already full -> clone it. + it.item.encoded = ( + self._packed_to_full(it, it.recv_main, self.hidden) + if any_sharded + else it.recv_main.clone() + ) + it.item.encoded_deepstack = None + it.pool.release_after_copy(it.pool_slot, it.item.encoded) + else: + # Legacy path. When sharded, the packed buffers are smaller than + # the image and cannot alias item.encoded, so build separate full + # tensors; when not sharded, hand the (full) buffers over directly. + if any_sharded: + it.item.encoded = self._packed_to_full( + it, it.recv_main, self.hidden + ) + it.item.encoded_deepstack = ( + self._packed_to_full( + it, it.recv_deepstack, self.hidden * self.num_deepstack + ) + if it.recv_deepstack is not None + else None + ) + else: + it.item.encoded = it.recv_main + it.item.encoded_deepstack = it.recv_deepstack + pairs = [(it.recv_main, it.recv_main.data_ptr())] + if it.recv_deepstack is not None: + pairs.append((it.recv_deepstack, it.recv_deepstack.data_ptr())) + _lazy_deregister(self.manager.engine, pairs) + # Receive concluded: release each room's bookkeeping from the prefill manager. + self._clear_receivers() + self._status = DONE + return DONE + + def _packed_to_full( + self, it: "_ItemReceive", packed: torch.Tensor, width: int + ) -> torch.Tensor: + """Scatter a PACKED shard buffer (this rank's rows, image-contiguous) into + a full-layout ``[n_tokens, width]`` tensor, placing each image's shard at + its absolute row offset and leaving the non-owned rows for ``reassemble`` + to fill. The packed buffer holds ``sum(row_counts)`` rows; the full tensor + holds ``sum(spans)`` rows (== n_tokens). Identity images (row_count==span, + row_start==0) copy whole, so a buffer with no real sharding round-trips + unchanged.""" + full = torch.empty( + (it.n_tokens, width), dtype=packed.dtype, device=packed.device + ) + packed_cursor = 0 + full_cursor = 0 + for span, row_start, row_count in zip(it.spans, it.row_starts, it.row_counts): + if row_count > 0: + full[full_cursor + row_start : full_cursor + row_start + row_count] = ( + packed[packed_cursor : packed_cursor + row_count] + ) + packed_cursor += row_count + full_cursor += span + return full + + def reassemble(self, nccl_group: Any, group_ranks: Sequence[int]) -> None: + """Rebuild full embeddings from row shards via per-image broadcasts. + + Must be called RANK-LOCKSTEP on every attn-TP rank, only after the + rank-agreed DONE (the drain's MIN all-reduce) and BEFORE any forward + consumes ``item.encoded``: until then each rank's buffer holds only its + own shard rows, the rest is uninitialized memory. All ranks iterate the + identical items/images in identical order issuing identical collectives, + which also requires the caller to run on the NON-overlap event loop (a + second thread launching forward collectives concurrently would break the + cross-rank launch-order guarantee NCCL needs across communicators). + + ``group_ranks`` is the attn-TP group as GLOBAL ranks + (``mapping.attn.tp_group``): ``dist.broadcast`` takes a global src rank, + and group rank p == global rank p only in the no-DP single-engine case. + Per-image sub-range broadcasts (2 x shard_size x n_images launches per + request). + + No-op for identity images and when sharding is off; safe after _fail + (items were dropped). Idempotence is NOT required: called exactly once + per admitted request by the drain. + + Broadcasts target ``item.encoded`` (the PUBLISHED tensor), never + ``recv_main``: on the pooled path the publish step cloned the rows out + and queued the slot for reuse after the clone completes, so ``recv_main`` + may later belong to the next tenant; on the legacy path ``item.encoded`` + IS ``recv_main``, so the two are equivalent there. + """ + if self._shard_size <= 1 or self._status is not DONE: + return + for it in self._items: + main = it.item.encoded + deep = it.item.encoded_deepstack + cursor = 0 + for idx, span in enumerate(it.spans): + if not it.sharded[idx]: + cursor += span + continue + for p in range(self._shard_size): + start, count = shard_rows(span, p, self._shard_size) + if count == 0: + continue + src = group_ranks[p] + dist.broadcast( + main[cursor + start : cursor + start + count], + src=src, + group=nccl_group, + ) + if deep is not None: + dist.broadcast( + deep[cursor + start : cursor + start + count], + src=src, + group=nccl_group, + ) + cursor += span + + def _clear_receivers(self) -> None: + """Release each receiver's per-room manager bookkeeping (terminal paths only; + no-op for fake receivers without clear()).""" + for it in self._items: + for receiver in it.receivers: + clear = getattr(receiver, "clear", None) + if clear is not None: + clear() + + def _fail(self) -> None: + """Tear down on failure and drop our references. A SIBLING image's + write may still be in flight into the item buffer, so reclamation is + deferred on both paths: pooled slots go to quarantine (under the + lifetime MR a late write would land SILENTLY in the next tenant's + data), legacy buffers go to the lazy-deregistration queue. + ``item.encoded`` is left unset -- the request is being failed, not + served. + """ + for it in self._items: + if it.pool is not None: + it.pool.quarantine(it.pool_slot, _RECV_POOL_QUARANTINE_S) + else: + pairs = [(it.recv_main, it.recv_main.data_ptr())] + if it.recv_deepstack is not None: + pairs.append((it.recv_deepstack, it.recv_deepstack.data_ptr())) + _lazy_deregister(self.manager.engine, pairs) + self._clear_receivers() + self._items = [] + self._status = FAILED + + @property + def status(self) -> str: + return self._status + + def release(self) -> None: + """Best-effort teardown for an externally-driven abort. + + Used when a rank-agreed FAILED is reached but THIS rank had not yet seen a + local Failed poll (so its buffers are still registered): free/deregister + them so a reused allocator address is never left double-registered. + Idempotent and a no-op once terminal (poll() already tore a DONE/FAILED job + down). After release() the job is FAILED. + """ + if self._status is not PENDING: + return + self._fail() + + +def start_embedding_receive( + items: Sequence[MultimodalDataItem], + manager: Any, + *, + hidden: int, + num_deepstack: int, + dtype: torch.dtype, + device: torch.device | str, + receiver_factory: ReceiverFactory | None = None, + shard_rank: int = 0, + shard_size: int = 1, +) -> EmbeddingReceiveJob: + """Begin (non-blocking) the per-item EPD embedding receive for one request. + + See :class:`EmbeddingReceiveJob`. The handshake for each EPD-routed item + rides on ``item.encode_handshake`` (a dict ``{bootstrap_room, bootstrap_host, + bootstrap_port}``); items without one are left to the vision tower. Returns a + job whose ``poll()`` the caller drives every cycle until DONE/FAILED; the + request must not be admitted to a prefill forward before then. With + ``shard_size > 1`` the caller must also run ``job.reassemble()`` rank- + lockstep after the rank-agreed DONE, before the request's first forward. + """ + return EmbeddingReceiveJob( + items, + manager, + hidden=hidden, + num_deepstack=num_deepstack, + dtype=dtype, + device=device, + receiver_factory=receiver_factory, + shard_rank=shard_rank, + shard_size=shard_size, + ) + + +def receive_encoded_embeddings( + items: Sequence[MultimodalDataItem], + manager: Any, + *, + hidden: int, + num_deepstack: int, + dtype: torch.dtype, + device: torch.device | str, + timeout: float = 60.0, + receiver_factory: ReceiverFactory | None = None, +) -> None: + """BLOCKING wrapper: fill ``item.encoded`` from the per-item transfers. + + For synchronous callers and the CPU unit tests: the poll-driven + :class:`EmbeddingReceiveJob` spun to completion in place (start, then + busy-poll until DONE, raising on FAILED or timeout). Event-loop code should + drive ``poll()`` once per cycle instead. + + The handshake for each EPD-routed item rides on ``item.encode_handshake`` + (``{bootstrap_room, bootstrap_host, bootstrap_port}``). ``hidden`` is the main + embedding width, ``num_deepstack`` the deepstack level count (0 if absent), + and ``dtype`` must match the encode worker's tower output dtype (asserted on + the encode side before the unchecked RDMA write). + """ + job = start_embedding_receive( + items, + manager, + hidden=hidden, + num_deepstack=num_deepstack, + dtype=dtype, + device=device, + receiver_factory=receiver_factory, + ) + deadline = time.monotonic() + timeout + while True: + status = job.poll() + if status == DONE: + return + if status == FAILED: + raise RuntimeError("encode->prefill embedding transfer failed") + if time.monotonic() > deadline: + raise TimeoutError( + f"embedding receive timed out after {timeout:.1f}s " + f"(job still {status})" + ) + time.sleep(0.0005) + + +def build_prefill_embedding_manager(server_args, global_rank, is_multimodal_active): + """Stand up the EPD encode->prefill embedding sink, if applicable. + + Only a multimodal *prefill* node receives embeddings from encode workers. Built + as one TP group (the embedding transport allows prefill_tp = any multiple of + encode_tp) with no fixed base buffer (receive buffers are allocated per image + at receive time). Construction spawns a daemon status thread, so build it + exactly once per rank. Returns None for decode/encode/text-only nodes. + """ + if server_args.disaggregation_mode != "prefill" or not is_multimodal_active: + return None + + from tokenspeed.runtime.pd.epd.embedding_transfer import ( + MooncakeEmbeddingManagerPrefill, + ) + from tokenspeed.runtime.pd.epd.entities import EmbeddingArgs, EmbeddingManagerArgs + + emb_mgr_args = EmbeddingManagerArgs( + bootstrap_port=server_args.disaggregation_bootstrap_port, + tp_size=server_args.mapping.attn.tp_size, + ) + emb_args = EmbeddingArgs( + engine_rank=global_rank, + gpu_id=global_rank, + ib_device=server_args.disaggregation_ib_device, + embedding_data_ptr=0, + embedding_data_len=0, + ) + return MooncakeEmbeddingManagerPrefill(emb_mgr_args, emb_args) + + +class EpdPrefillAdmission: + """EPD prefill-side embedding receive + rank-synced admission. + + Owns the encode->prefill embedding sink (MooncakeEmbeddingManagerPrefill), the + set of requests whose per-image embeddings are still arriving (``_pending``), + and the optional NCCL row-shard reassembly. Exists ONLY on a multimodal prefill + node; the EventLoop holds one (or None) and drives it each non-overlap cycle. + + DECIDE/ACT split: ``drain()`` polls the staged receives, applies the rank- + lockstep MIN all-reduce + reassembly, and RETURNS ``(admitted, failed)`` + decisions. The EventLoop performs the acts those imply (P->D sender + register/abort, scheduler submit, output-processor finish) -- they touch + EventLoop collaborators, so they stay there. + """ + + def __init__( + self, + *, + manager, + device, + hidden, + num_deepstack, + dtype, + attn_tp_rank, + attn_tp_size, + attn_tp_cpu_group, + attn_tp_group, + pg_manager, + ): + self._manager = manager + self._device = device + self._hidden = hidden + self._num_deepstack = num_deepstack + self._dtype = dtype + self._attn_tp_rank = attn_tp_rank + self._attn_tp_size = attn_tp_size + self._attn_tp_cpu_group = attn_tp_cpu_group + # Requests whose per-image encode->prefill embeddings are still arriving: + # registered but NOT yet submitted to the scheduler; drain() polls them each + # cycle and admits (rank-synced) only once ready. Rank-identical in + # length+order across attn-TP ranks (recv_reqs broadcasts the new-request + # set), which drain()'s MIN all-reduce relies on. + self._pending: list = [] + # Deadline for an EPD request's per-image embeddings to all arrive; past it + # the request is aborted (not waited on forever). Reuse the PD KV-receive + # wait knob (default 300s): the prefill waiting on the encode->prefill + # embedding transfer is the direct analog of the decode waiting on the + # prefill->decode KV transfer, so one operator knob covers both. + self._embed_timeout: float = float( + envs.TOKENSPEED_DISAGGREGATION_WAITING_TIMEOUT.get() + ) + + # EPD embedding row-sharding: each attn-TP rank receives only 1/N of every + # image's rows over the wire; the full embedding is rebuilt by an NCCL + # all-gather in drain() (job.reassemble), which runs on the prefill's + # NON-overlap loop so the drain and the forward launch from one thread in + # the same order on every rank -- the cross-rank launch-order consistency + # NCCL requires across communicators. + self._shard_embeddings = False + self._nccl_group = None + self._group_ranks = tuple(attn_tp_group) + shard_flag = False + if attn_tp_size > 1: + shard_flag = bool(envs.TOKENSPEED_EPD_EMBEDDING_SHARD.get()) + # The flag is a per-process env read but gates a GROUP collective: + # torn across ranks (e.g. set on one node of a multi-node prefill), + # flag-on ranks would join the warmup broadcast below while flag-off + # ranks never do -- a silent boot hang. Agree first, loud. + flag_t = torch.tensor( + [int(shard_flag), -int(shard_flag)], dtype=torch.int32 + ) + dist.all_reduce(flag_t, op=dist.ReduceOp.MIN, group=attn_tp_cpu_group) + if flag_t[0].item() != -flag_t[1].item(): + raise RuntimeError( + "TOKENSPEED_EPD_EMBEDDING_SHARD differs across attn-TP ranks; " + "set it identically on every node of the prefill engine" + ) + if shard_flag: + self._nccl_group = pg_manager.get_process_group("nccl", attn_tp_group) + self._shard_embeddings = True + # Warm the communicator at startup: NCCL initializes lazily on the + # first collective, and nothing else issues torch.distributed NCCL ops + # on this group in the scheduler process -- without this the FIRST + # admitted EPD request pays communicator init (hundreds of ms, all + # ranks) inside the drain, and a misconfigured group would surface on a + # customer request instead of at boot. + warmup = torch.zeros(1, device=device) + dist.broadcast(warmup, src=self._group_ranks[0], group=self._nccl_group) + torch.cuda.current_stream().synchronize() + logger.info( + "EPD embedding row-sharding enabled (attn_tp=%d, NCCL group warm)", + attn_tp_size, + ) + + def stage(self, request_id, mm_items) -> None: + """Begin the non-blocking per-image embedding receive and stage it by + request_id until its embeddings land (polled in drain()). The request payload + (spec/state/bootstrap) stays with the caller, keyed by request_id; this + controller tracks only the receive job (mirrors kv_transfer).""" + job = start_embedding_receive( + items=mm_items, + manager=self._manager, + hidden=self._hidden, + num_deepstack=self._num_deepstack, + dtype=self._dtype, + device=self._device, + shard_rank=self._attn_tp_rank, + shard_size=self._attn_tp_size if self._shard_embeddings else 1, + ) + self._pending.append((request_id, job, time.time())) + + def has_pending(self) -> bool: + return bool(self._pending) + + def drain(self): + """Poll staged EPD embedding receives; return ``(admitted_ids, failed_ids)``. + + poll/timeout/MIN-all-reduce/reassemble/release + ``_pending`` bookkeeping + happen here, rank-lockstep, keyed by request_id. The caller maps the ids back + to its staged request payloads and performs the acts (kv_transfer + register/abort, scheduler submit, output-processor finish). + + - admitted_ids: DONE on every rank, reassembled. + - failed_ids: FAILED/timed-out, job released. + """ + if not self._pending: + return [], [] # rank-identical emptiness -> all ranks skip the collective + + code_of = {FAILED: 0, PENDING: 1, DONE: 2} + codes = [code_of[job.poll()] for (_rid, job, _ts) in self._pending] + + # Timeout: a still-PENDING request whose per-image embeddings have not all + # arrived within the deadline is marked FAILED (-> rank-agreed abort below). + # Without this the prefill waits FOREVER if an embedding is ever lost (a + # degraded/dead encode worker, a network drop). Folded into the SAME MIN + # all-reduce: a timed-out job -> code 0 -> all ranks abort it together; + # union-of-timeout is rank-safe even if ranks cross the deadline a cycle + # apart (any rank's 0 propagates via MIN). + _now = time.time() + for _i in range(len(self._pending)): + if codes[_i] == 1 and (_now - self._pending[_i][2]) > self._embed_timeout: + codes[_i] = 0 + logger.warning( + "EPD embedding receive timed out after %.0fs for rid=%s; aborting", + self._embed_timeout, + self._pending[_i][0], + ) + + if self._attn_tp_size > 1: + t = torch.tensor(codes, dtype=torch.uint8, device="cpu") + dist.all_reduce(t, op=dist.ReduceOp.MIN, group=self._attn_tp_cpu_group) + codes = t.tolist() + + admitted = [] + failed = [] + leftover = [] + for (request_id, job, start_ts), code in zip(self._pending, codes): + if code == 2: # DONE on every rank + # Sharded receive: rebuild the full rows from the per-rank shards + # FIRST -- item.encoded is shard-only until this runs. Rank-lockstep- + # safe: codes are identical post-MIN and _pending is rank-identical + # in length/order, so every rank issues identical collectives in + # identical order this cycle. + if self._shard_embeddings: + job.reassemble(self._nccl_group, self._group_ranks) + admitted.append(request_id) + elif code == 0: # FAILED/timed-out on some rank -> abort everywhere + job.release() + failed.append(request_id) + else: # still pending on some rank + leftover.append((request_id, job, start_ts)) + self._pending = leftover + return admitted, failed + + +def make_epd_prefill_admission( + server_args, + global_rank, + *, + model_config, + model_executor, + mapping, + attn_tp_rank, + attn_tp_size, + attn_tp_cpu_group, + pg_manager, +): + """Build the EPD-prefill admission controller, or None for non-EPD nodes. + + Returns None unless this is a multimodal *prefill* node (the only node that + receives encode->prefill embeddings).""" + manager = build_prefill_embedding_manager( + server_args, global_rank, model_config.is_multimodal_active + ) + if manager is None: + return None + # Extract the narrow model facts the controller needs (vision dtype, hidden + # width, deepstack width, device) here -- the controller holds these, not the + # whole model_executor (mirrors how create_kv_transfer takes a kv_args struct). + model = model_executor.model_runner.model + return EpdPrefillAdmission( + manager=manager, + device=model_executor.device, + hidden=model.config.hidden_size, + num_deepstack=getattr(model, "num_deepstack_embeddings", 0), + dtype=(getattr(model, "visual", None) or model.vision_tower).dtype, + attn_tp_rank=attn_tp_rank, + attn_tp_size=attn_tp_size, + attn_tp_cpu_group=attn_tp_cpu_group, + attn_tp_group=mapping.attn.tp_group, + pg_manager=pg_manager, + ) diff --git a/python/tokenspeed/runtime/pd/factory.py b/python/tokenspeed/runtime/pd/factory.py new file mode 100644 index 0000000..5b8d1ca --- /dev/null +++ b/python/tokenspeed/runtime/pd/factory.py @@ -0,0 +1,140 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Factories for PD KV transfer helpers.""" + +from tokenspeed.runtime.pd.decode_executor import DisaggDecodeExecutor +from tokenspeed.runtime.pd.mooncake.entities import KVArgs, KVManagerArgs +from tokenspeed.runtime.pd.prefill_executor import DisaggPrefillExecutor +from tokenspeed.runtime.pd.utils import TransferBackend + + +def _get_contiguous_buf_unit_lens(pool, item_lens): + getter = getattr(pool, "get_contiguous_buf_unit_lens", None) + if getter is None: + return [1] * len(item_lens) + unit_lens = list(getter()) + if len(unit_lens) != len(item_lens): + raise ValueError( + f"contiguous buffer unit count mismatch: units={len(unit_lens)}, items={len(item_lens)}" + ) + return unit_lens + + +def get_kv_args( + engine_rank: int, + gpu_id, + ib_device, + token_to_kv_pool, + draft_token_to_kv_pool, + mamba_pool=None, +): + kv_data_ptrs, kv_data_lens, kv_item_lens = ( + token_to_kv_pool.get_contiguous_buf_infos() + ) + kv_unit_lens = _get_contiguous_buf_unit_lens(token_to_kv_pool, kv_item_lens) + # [[layer0buf0, layer0buf1...], [layer1buf0, layer1buf1...], ...] + offsets = token_to_kv_pool.get_layerwise_buf_info_offsets() + target_layer_num = token_to_kv_pool.layer_num + kv_layer_ids = list(getattr(token_to_kv_pool, "layer_ids", range(target_layer_num))) + + if draft_token_to_kv_pool is not None: + draft_layer_num = draft_token_to_kv_pool.layer_num + # 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 = ( + draft_token_to_kv_pool.get_contiguous_buf_infos() + ) + draft_kv_unit_lens = _get_contiguous_buf_unit_lens( + draft_token_to_kv_pool, draft_kv_item_lens + ) + draft_offsets = draft_token_to_kv_pool.get_layerwise_buf_info_offsets( + len(kv_data_ptrs) + ) + + kv_data_ptrs += draft_kv_data_ptrs + kv_data_lens += draft_kv_data_lens + kv_item_lens += draft_kv_item_lens + kv_unit_lens += draft_kv_unit_lens + offsets += draft_offsets + draft_base_layer_id = ( + max(kv_layer_ids) + 1 if kv_layer_ids else target_layer_num + ) + kv_layer_ids += list( + range(draft_base_layer_id, draft_base_layer_id + draft_layer_num) + ) + else: + draft_layer_num = 0 + + state_data_ptrs = [] + state_data_lens = [] + state_item_lens = [] + state_unit_lens = [] + state_type = "none" + state_layer_ids = [] + if mamba_pool is not None: + state_data_ptrs, state_data_lens, state_item_lens = ( + mamba_pool.get_contiguous_buf_infos() + ) + state_unit_lens = _get_contiguous_buf_unit_lens(mamba_pool, state_item_lens) + state_layer_ids = mamba_pool.get_contiguous_buf_layer_ids() + state_type = "mamba" + + kv_args = KVArgs( + engine_rank=engine_rank, + kv_data_ptrs=kv_data_ptrs, + kv_data_lens=kv_data_lens, + kv_item_lens=kv_item_lens, + target_layer_num=target_layer_num, + draft_layer_num=draft_layer_num, + kv_layer_ids=kv_layer_ids, + kv_unit_lens=kv_unit_lens, + state_data_ptrs=state_data_ptrs, + state_data_lens=state_data_lens, + state_item_lens=state_item_lens, + state_unit_lens=state_unit_lens, + state_type=state_type, + state_layer_ids=state_layer_ids, + mamba_offsets=[], + offsets=offsets, + aux_data_ptrs=[], + aux_data_lens=[], + aux_item_lens=[], + ib_device=ib_device, + gpu_id=gpu_id, + ) + + return kv_args + + +def create_kv_transfer( + mode: str, + backend: TransferBackend, + args: KVManagerArgs, + kv_args: KVArgs, + gloo_group, + page_size, +): + if mode == "prefill": + return DisaggPrefillExecutor(backend, args, kv_args, gloo_group, page_size) + elif mode == "decode": + return DisaggDecodeExecutor(backend, args, kv_args, gloo_group, page_size) + else: + raise NotImplementedError(f"Unsupported disaggregation mode: {mode}") diff --git a/python/tokenspeed/runtime/pd/kv_events.py b/python/tokenspeed/runtime/pd/kv_events.py new file mode 100755 index 0000000..d0d25df --- /dev/null +++ b/python/tokenspeed/runtime/pd/kv_events.py @@ -0,0 +1,479 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""KV cache event definitions and publisher implementations.""" + +import atexit +import logging +import queue +import threading +import time +from abc import ABC, abstractmethod +from collections import deque +from collections.abc import Callable, Iterable +from itertools import count +from queue import Queue +from typing import Any + +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: int | None = 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: int | None + token_ids: list[int] + block_size: int + + +class BlockRemoved(KVCacheEvent): + block_hashes: list[int] + + +class AllBlocksCleared(KVCacheEvent): + pass + + +class KVEventBatch(EventBatch): + events: list[BlockStored | BlockRemoved | AllBlocksCleared] + + +def scheduler_kv_event_to_wire_event( + event: Any, +) -> BlockStored | BlockRemoved: + """Translate a scheduler-native KV event into the ZMQ wire struct.""" + kind = event.kind + if kind == "BlockStored": + return BlockStored( + block_hashes=[int(block_hash) for block_hash in event.block_hashes], + parent_block_hash=( + None + if event.parent_block_hash is None + else int(event.parent_block_hash) + ), + token_ids=[int(token_id) for token_id in event.token_ids], + block_size=int(event.block_size), + ) + if kind == "BlockRemoved": + return BlockRemoved( + block_hashes=[int(block_hash) for block_hash in event.block_hashes] + ) + raise TypeError(f"Unsupported scheduler KV event kind: {kind}") + + +def scheduler_kv_events_to_wire_events( + events: Iterable[Any], +) -> list[BlockStored | BlockRemoved]: + return [scheduler_kv_event_to_wire_event(event) for event in events] + + +def drain_scheduler_kv_events(scheduler: Any, enabled: bool) -> list[Any]: + if not enabled: + return [] + + drain_kv_events = getattr(scheduler, "drain_kv_events", None) + if drain_kv_events is None: + raise RuntimeError( + "KV cache events require a tokenspeed_scheduler extension with " + "Scheduler.drain_kv_events() support." + ) + + return list(drain_kv_events()) + + +class EventPublisher(ABC): + """Lightweight publisher for ``EventBatch`` batches with 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 + this implementation: + + - Each DP rank has its own ``EventPublisher``. + - Publishers annotate events with the DP rank. + - Consumers can 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: EventBatch) -> 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: str | None = 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[EventBatch | None](maxsize=max_queue_size) + self._buffer = deque[tuple[int, bytes]](maxlen=buffer_steps) + + # ZMQ sockets + self._ctx = zmq.Context.instance() + self._pub: zmq.Socket | None = None + self._replay: zmq.Socket | None = 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() + + if self._pub is None: + raise RuntimeError("Publisher socket is not initialized.") + + 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 exc: + logger.exception("Error in replay: %s", exc) + + # --- 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 exc: + # Publishing failed; back-off a bit to avoid a tight error loop + logger.exception("Error in publisher thread: %s", exc) + time.sleep(0.1) + + def _service_replay(self) -> None: + """If a replay request is waiting, send buffered batches.""" + if self._replay is None: + raise RuntimeError("Replay socket is not initialized.") + + 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: str | None, data_parallel_rank: int + ) -> str | None: + """Offset the port in an endpoint by the data parallel rank. + + Args: + endpoint: The endpoint string, for example ``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 a + 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 | None = None + """The publisher to use for publishing kv events. Can be "null", "zmq". + Defaults to "zmq" when events are enabled and "null" otherwise. + """ + + endpoint: str = "tcp://*:5557" + """The zmq endpoint to use for publishing kv events. + """ + + replay_endpoint: str | None = 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. + """ + + enable_kv_cache_events: bool = False + """Whether to publish scheduler KV cache mutation 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: str | None, attn_dp_rank: int = 0) -> EventPublisher: + """Create publisher from a config mapping.""" + if not config: + return NullEventPublisher(attn_dp_rank=attn_dp_rank) + config = KVEventsConfig.from_cli(config) + config_dict = config.model_dump() + enabled = bool(config_dict.pop("enable_kv_cache_events", False)) + if not enabled: + return NullEventPublisher(attn_dp_rank=attn_dp_rank) + + kind = config_dict.pop("publisher", None) + if kind is None: + kind = "zmq" if enabled else "null" + if kind == "null": + return NullEventPublisher(attn_dp_rank=attn_dp_rank) + 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) + + @classmethod + def is_enabled(cls, config: str | None) -> bool: + if not config: + return False + return KVEventsConfig.from_cli(config).enable_kv_cache_events diff --git a/python/tokenspeed/runtime/pd/mooncake/__init__.py b/python/tokenspeed/runtime/pd/mooncake/__init__.py new file mode 100644 index 0000000..92d95f8 --- /dev/null +++ b/python/tokenspeed/runtime/pd/mooncake/__init__.py @@ -0,0 +1,41 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from tokenspeed.runtime.pd.mooncake.conn import MooncakeKVBootstrapServer +from tokenspeed.runtime.pd.mooncake.decode import ( + MooncakeKVManagerDecode, +) +from tokenspeed.runtime.pd.mooncake.prefill import ( + MooncakeKVManagerPrefill, +) +from tokenspeed.runtime.pd.mooncake.receiver import ( + MooncakeKVReceiver, +) +from tokenspeed.runtime.pd.mooncake.sender import ( + MooncakeKVSender, +) + +__all__ = ( + "MooncakeKVBootstrapServer", + "MooncakeKVManagerDecode", + "MooncakeKVManagerPrefill", + "MooncakeKVReceiver", + "MooncakeKVSender", +) diff --git a/python/tokenspeed/runtime/pd/mooncake/async_conn.py b/python/tokenspeed/runtime/pd/mooncake/async_conn.py new file mode 100644 index 0000000..c1bc842 --- /dev/null +++ b/python/tokenspeed/runtime/pd/mooncake/async_conn.py @@ -0,0 +1,566 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import dataclasses +import threading +import time +from contextlib import contextmanager +from typing import Any + +import numpy as np +import numpy.typing as npt + +from tokenspeed.runtime.pd.base.status import TransferPoll +from tokenspeed.runtime.pd.mooncake.entities import ( + KVArgs, + TransferInfo, + TransferKVChunk, +) +from tokenspeed.runtime.pd.mooncake.prefill import ( + MooncakeKVManagerPrefill as MooncakeKVManager, +) +from tokenspeed.runtime.pd.utils import ( + DisaggregationMode, + FastQueue, + StepCounter, + group_concurrent_contiguous, +) +from tokenspeed.runtime.utils import get_colorful_logger +from tokenspeed.runtime.utils.server_args import ServerArgs + +logger = get_colorful_logger(__name__) + + +@dataclasses.dataclass +class WriteRequest: + trans_info: TransferInfo + dst_ranks_info: tuple[str, int, int] + prefill_kv_blocks: npt.NDArray[np.int64] + dst_kv_blocks: npt.NDArray[np.int64] + submit_bids: list[int] = dataclasses.field(default_factory=list) + + +@dataclasses.dataclass +class LayerWiseTask: + kv_chunk: TransferKVChunk + begin_cache_step: int + aux_step: int + next_layer_id: int = 0 + write_requests: list[WriteRequest] = dataclasses.field(default_factory=list) + polls: list[bool] = dataclasses.field(default_factory=list) + + +class MooncakeAsyncKVManager(MooncakeKVManager): + def __init__( + self, + args: KVArgs, + disaggregation_mode: DisaggregationMode, + server_args: ServerArgs, + is_mla_backend: bool | None = False, + draft_is_mla_backend: bool | None = False, + ): + super().__init__( + args, disaggregation_mode, server_args, is_mla_backend, draft_is_mla_backend + ) + self.target_layer_num = self.kv_args.target_layer_num + self.draft_layer_num = self.kv_args.draft_layer_num + self.layer_num = self.target_layer_num + self.draft_layer_num + self.is_mla_backend = is_mla_backend + self.draft_is_mla_backend = draft_is_mla_backend + self.kv_cache_quant_method = server_args.kv_cache_quant_method + self.submit_interval = server_args.disaggregation_layerwise_interval + if self.submit_interval <= 0: + raise ValueError("submit_interval must be positive") + self.current_transfer_batch: list[tuple[TransferKVChunk, int]] = [] + self.offsets = args.offsets + + def _init_offsets(self): + if self.is_mla_backend: + if self.kv_cache_quant_method == "per_token_head": + target_offsets = [ + [i * self.target_layer_num + layer_id for i in range(3)] + for layer_id in range(self.target_layer_num) + ] + draft_offsets = [ + [3 * self.target_layer_num + layer_id] + for layer_id in range(self.draft_layer_num) + ] + else: + target_offsets = [ + [layer_id] for layer_id in range(self.target_layer_num) + ] + draft_offsets = [ + [self.target_layer_num + layer_id] + for layer_id in range(self.draft_layer_num) + ] + else: + target_offsets = [ + [i * self.layer_num + layer_id for i in range(2)] + for layer_id in range(self.target_layer_num) + ] + draft_offsets = [ + [ + i * self.layer_num + self.target_layer_num + layer_id + for i in range(2) + ] + for layer_id in range(self.draft_layer_num) + ] + self.offsets = target_offsets + draft_offsets + + def start_transfer_thread( + self, transfer_thread_pool_size: int, transfer_queue_size: int + ): + self.transfer_queues: list[FastQueue] = [FastQueue()] + for queue in self.transfer_queues: + threading.Thread( + target=self.async_transfer_worker, args=(queue,), daemon=True + ).start() + + def register_step_counter(self, step_counter: StepCounter): + self.step_counter = step_counter + + @contextmanager + def add_batch(self, is_idle: bool): + yield # add transfer request + + if not is_idle: + begin_cache_step, begin_aux_step = self.step_counter.current_step() + for kv_chunk, shard_idx in self.current_transfer_batch: + self.transfer_queues[shard_idx].put( + LayerWiseTask( + kv_chunk=kv_chunk, + begin_cache_step=begin_cache_step, + aux_step=begin_aux_step if kv_chunk.is_last else None, + ) + ) + + # advance to step of next batch no matter if batch is empty + self.step_counter.advance_step( + delta_cache_step=self.layer_num, delta_aux_step=1 + ) + + self.current_transfer_batch.clear() + + def add_transfer_request( + self, + bootstrap_room: int, + kv_indices: npt.NDArray[np.int64], + index_slice: slice, + is_last: bool, + aux_index: int | None = None, + mla_l1_5_args: tuple[Any, Any] | None = None, + ): + logger.debug("async manager add_transfer_request") + if self.disaggregation_mode != DisaggregationMode.PREFILL: + raise RuntimeError("Transfer requests can only be added in prefill mode.") + if is_last and aux_index is None: + raise ValueError("aux_index must be set for the last transfer chunk.") + + if ( + bootstrap_room not in self.request_status + or self.check_status(bootstrap_room) == TransferPoll.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 + + # 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.split(":")[1]) for session in dst_infos) + shard_idx = session_port_sum % len(self.transfer_queues) + + kv_chunk = TransferKVChunk( + room=bootstrap_room, + prefill_kv_indices=kv_indices, + index_slice=index_slice, + is_last=is_last, + prefill_aux_index=aux_index, + mla_l1_5_args=mla_l1_5_args, + ) + self.current_transfer_batch.append((kv_chunk, shard_idx)) + + def submit_aux( + self, + mooncake_session_id: str, + prefill_aux_index: int, + dst_aux_ptrs: list[int], + dst_aux_index: int, + ): + # Submit transfer for all aux buffers (output_ids, logprobs, cached_tokens, etc.) + batch_ids = [] + for aux_data_ptr, aux_item_len, dst_aux_ptr in zip( + self.kv_args.aux_data_ptrs, self.kv_args.aux_item_lens, dst_aux_ptrs + ): + prefill_aux_addr = aux_data_ptr + prefill_aux_index * aux_item_len + decode_aux_addr = dst_aux_ptr + dst_aux_index * aux_item_len + bid = self.engine.transfer_submit_write( + mooncake_session_id, prefill_aux_addr, decode_aux_addr, aux_item_len + ) + batch_ids.append(bid) + # The public interface returns a single batch id, so expose the last + # submitted id even though this helper may issue multiple writes. + return batch_ids[-1] if batch_ids else -1 + + 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 submit_layer_cache( + self, + mooncake_session_id: str, + begin_layer_id: int, # [begin_layer_id, end_layer_id) + end_layer_id: int, + prefill_kv_blocks: npt.NDArray[np.int64], + dst_kv_blocks: npt.NDArray[np.int64], + ) -> int: + dst_kv_ptrs = self.decode_kv_args_table[mooncake_session_id].dst_kv_ptrs + transfer_blocks = [] + + def submit_one_cache(ptr_offset): + src_ptr = self.kv_args.kv_data_ptrs[ptr_offset] + dst_ptr = dst_kv_ptrs[ptr_offset] + item_len = self.kv_args.kv_item_lens[ptr_offset] + 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)) + + for layer_id in range(begin_layer_id, end_layer_id): + for offset in self.offsets[layer_id]: + submit_one_cache(offset) + return self._transfer_data(mooncake_session_id, transfer_blocks) + + def async_transfer_worker(self, transfer_queue: FastQueue): + def discard_finished_bid_inplace(submit_bids: list[int]): + finished_cnt = 0 + failed = False + for bid in submit_bids: + status = self.engine.transfer_check_status(bid) + if status == 1: + finished_cnt += 1 + elif status == -1: + failed = True + if self.kv_transfer_metrics: + self.kv_transfer_metrics.record_kv_transfer_timeout() + logger.error("Transfer timeout detected!") + break + elif status == -2: + failed = True + if self.kv_transfer_metrics: + self.kv_transfer_metrics.record_kv_transfer_failure() + logger.error("Transfer failed detected") + break + else: + failed = status != 0 + break + submit_bids[:] = submit_bids[finished_cnt:] + return not failed + + def discard_tasks( + tasks: list[LayerWiseTask], dropped: list[LayerWiseTask] + ) -> list[LayerWiseTask]: + dropped_rooms = set(id(task) for task in dropped) + return [task for task in tasks if id(task) not in dropped_rooms] + + def query_ready_step(task: LayerWiseTask) -> tuple[int, int]: + if task.next_layer_id < self.layer_num: + ready_cache_step = self.step_counter.query_ready_cache_step() + if not StepCounter.is_step_ready( + ready_cache_step, task.begin_cache_step + task.next_layer_id + ): + time.sleep(1e-3) + elif task.kv_chunk.is_last: + ready_aux_step = self.step_counter.query_ready_aux_step() + if not StepCounter.is_step_ready(ready_aux_step, task.aux_step): + time.sleep(1e-3) + + ready_cache_step = self.step_counter.query_ready_cache_step() + ready_aux_step = self.step_counter.query_ready_aux_step() + return ready_cache_step, ready_aux_step + + def get_new_tasks(blocking: bool) -> list[LayerWiseTask]: + new_tasks: list[LayerWiseTask] = [] + if blocking: + new_tasks.append(transfer_queue.get()) + while True: + try: + new_tasks.append(transfer_queue.get_nowait()) + except FastQueue.Empty: + break + + return new_tasks + + def initialize(tasks: list[LayerWiseTask]) -> list[LayerWiseTask]: + abort_tasks: list[LayerWiseTask] = [] + + for task in tasks: + kv_chunk = task.kv_chunk + reqs_to_be_processed = ( + self.transfer_infos[kv_chunk.room].values() + if kv_chunk.room in self.transfer_infos + else [] + ) + + for req in reqs_to_be_processed: + if req.is_dummy: + task.polls.append(True) + abort_tasks.append(task) + break + + with self.session_lock: + if self._is_session_failed(req.mooncake_session_id): + logger.info( + "Blocked transfer due to failed session (room=%s, session=%s).", + kv_chunk.room, + req.mooncake_session_id, + ) + self.record_failure( + kv_chunk.room, + f"Decode instance could be dead, remote mooncake session {req.mooncake_session_id} is not alive", + ) + task.polls.append(False) + abort_tasks.append(task) + break + + resolved = self.resolve_transfer_indices(kv_chunk, req) + + # Group by indices + prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous( + resolved.src_indices, resolved.dst_indices + ) + task.write_requests.append( + WriteRequest( + trans_info=req, + dst_ranks_info=(req.endpoint, req.dst_port, req.room), + prefill_kv_blocks=prefill_kv_blocks, + dst_kv_blocks=dst_kv_blocks, + ) + ) + + return abort_tasks + + def submit_transfer( + tasks: list[LayerWiseTask], ready_cache_step: int, ready_aux_step: int + ) -> tuple[list[LayerWiseTask], list[LayerWiseTask]]: + abort_tasks: list[LayerWiseTask] = [] + complete_tasks: list[LayerWiseTask] = [] + + for task in tasks: + kv_chunk = task.kv_chunk + # submit layer cache + if task.next_layer_id < self.layer_num and StepCounter.is_step_ready( + ready_cache_step, task.begin_cache_step + task.next_layer_id + ): + for req in task.write_requests: + if ( + (task.next_layer_id + 1) % self.submit_interval == 0 + or task.next_layer_id == self.layer_num - 1 + ): + ret = self.submit_layer_cache( + req.trans_info.mooncake_session_id, + (task.next_layer_id // self.submit_interval) + * self.submit_interval, + task.next_layer_id + 1, + req.prefill_kv_blocks, + req.dst_kv_blocks, + ) + + if ret != 0: + if self.kv_transfer_metrics: + self.kv_transfer_metrics.record_kv_transfer_failure() + logger.error("Transfer failed detected!") + task.polls.append(False) + abort_tasks.append(task) + + with self.session_lock: + self.session_failures[ + req.trans_info.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.trans_info.mooncake_session_id + ] + >= 1 + ): + self._mark_session_failed( + req.trans_info.mooncake_session_id, + reason="submit_layer_cache", + ) + logger.error( + "Session %s failed.", + req.trans_info.mooncake_session_id, + ) + self.record_failure( + kv_chunk.room, + f"Failed to send kv chunk of {kv_chunk.room} to {req.trans_info.endpoint}:{req.trans_info.dst_port}", + ) + break + + task.next_layer_id += 1 + + # submit aux data + if ( + kv_chunk.is_last + and task.aux_step is not None + and StepCounter.is_step_ready(ready_aux_step, task.aux_step) + ): + task.aux_step = None # reset to None to mark aux has been submitted + if kv_chunk.is_last: + for req in task.write_requests: + aux_bid = self.submit_aux( + req.trans_info.mooncake_session_id, + kv_chunk.prefill_aux_index, + self.decode_kv_args_table[ + req.trans_info.mooncake_session_id + ].dst_aux_ptrs, + req.trans_info.dst_aux_index, + ) + req.submit_bids.append(aux_bid) + + if task.next_layer_id == self.layer_num and task.aux_step is None: + complete_tasks.append(task) + + return complete_tasks, abort_tasks + + def pop_transferred(tasks: list[LayerWiseTask]) -> list[LayerWiseTask]: + complete_tasks: list[LayerWiseTask] = [] + for task in tasks: + kv_chunk = task.kv_chunk + for req in task.write_requests[ + len(task.polls) : + ]: # only check the uncompleted requests + success = discard_finished_bid_inplace(req.submit_bids) + if success: + if len(req.submit_bids) == 0: + task.polls.append(True) + else: + task.polls.append(False) + with self.session_lock: + self.session_failures[ + req.trans_info.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.trans_info.mooncake_session_id + ] + >= 1 + ): + self._mark_session_failed( + req.trans_info.mooncake_session_id, + reason="submit_status", + ) + logger.error( + "Session %s failed.", + req.trans_info.mooncake_session_id, + ) + self.record_failure( + kv_chunk.room, + f"Failed to send kv chunk of {kv_chunk.room} to {req.trans_info.endpoint}:{req.trans_info.dst_port}", + ) + break + + # all finished or any failed + if len(task.polls) == len(task.write_requests) or ( + len(task.polls) > 0 and not all(task.polls) + ): + complete_tasks.append(task) + + return complete_tasks + + def finalize(tasks: list[LayerWiseTask]) -> None: + for task in tasks: + kv_chunk = task.kv_chunk + status = ( + TransferPoll.Success if all(task.polls) else TransferPoll.Failed + ) + if ( + status == TransferPoll.Failed or kv_chunk.is_last + ): # last chunk or any failed + self.update_status(kv_chunk.room, status) + for packed_req in task.write_requests: + endpoint, dst_port, room = packed_req.dst_ranks_info + self.sync_status_to_decode_endpoint( + endpoint, dst_port, room, status, self.attn_tp_rank + ) + + if ( + kv_chunk.room not in self.request_status + or self.check_status(kv_chunk.room) == TransferPoll.Success + ): + if kv_chunk.room in self.transfer_infos: + self.transfer_infos.pop(kv_chunk.room) + + pending_tasks: list[LayerWiseTask] = [] + inflight_tasks: list[LayerWiseTask] = [] + while True: + try: + if ( + len( + new_tasks := get_new_tasks( + blocking=(len(pending_tasks) + len(inflight_tasks) == 0) + ) + ) + > 0 + ): + if len(abort_tasks := initialize(new_tasks)) > 0: + finalize(abort_tasks) + new_tasks = discard_tasks(new_tasks, abort_tasks) + + pending_tasks.extend(new_tasks) + + if len(pending_tasks) > 0: + ready_cache_step, ready_aux_step = query_ready_step( + pending_tasks[0] + ) # only wait the first task + submited_tasks, abort_tasks = submit_transfer( + pending_tasks, ready_cache_step, ready_aux_step + ) + finalize(abort_tasks) + pending_tasks = discard_tasks( + pending_tasks, submited_tasks + abort_tasks + ) + inflight_tasks.extend(submited_tasks) + + if len(complete_tasks := pop_transferred(inflight_tasks)) > 0: + finalize(complete_tasks) + inflight_tasks = discard_tasks(inflight_tasks, complete_tasks) + + except Exception as exc: + raise RuntimeError( + f"Transfer thread failed because of {exc}. Prefill instance " + f"with bootstrap_port={self.bootstrap_port} is dead." + ) from exc diff --git a/python/tokenspeed/runtime/pd/mooncake/conn.py b/python/tokenspeed/runtime/pd/mooncake/conn.py new file mode 100644 index 0000000..0b4e50c --- /dev/null +++ b/python/tokenspeed/runtime/pd/mooncake/conn.py @@ -0,0 +1,124 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from tokenspeed.runtime.metrics.collector import KVTransferMetrics +from tokenspeed.runtime.pd.base.bootstrap import DisaggBootstrapServerBase +from tokenspeed.runtime.pd.base.manager import DisaggManagerBase +from tokenspeed.runtime.pd.base.mooncake_engine import ( + MooncakeTransferEngine, +) +from tokenspeed.runtime.pd.mooncake.entities import KVArgs, KVManagerArgs +from tokenspeed.runtime.pd.utils import DisaggregationMode +from tokenspeed.runtime.utils.network import get_local_ip_by_remote + + +class MooncakeKVManagerBase(DisaggManagerBase): + """KV (prefill->decode) manager: the shared engine/socket/status FSM plus the + KV-specific args, MLA flags, dp-attention wiring, and metrics.""" + + def __init__( + self, + args: KVManagerArgs, + kv_args: KVArgs, + disaggregation_mode: DisaggregationMode, + ): + self.args = args + self.kv_args = kv_args + self.attn_tp_rank = args.attn_tp_rank + self.src_mode = "ON" if bool(args.enable_mla_l1_5_cache) else "OFF" + self.is_mla_backend = args.is_mla_backend + self.draft_is_mla_backend = args.draft_is_mla_backend + self.disaggregation_mode = disaggregation_mode + self.bootstrap_port = args.bootstrap_port + self.dist_init_addr = args.dist_init_addr + self.world_size = args.world_size + self.dp_size = args.dp_size + self.attn_dp_rank = args.attn_dp_rank + self.enable_dp_attention = args.enable_dp_attention + if not args.enable_dp_attention and args.dp_size != 1: + raise ValueError( + "If dp_attention is not enabled, dp size must be 1 in disaggregation mode." + ) + + if hasattr(args, "enable_metrics") and args.enable_metrics: + labels = { + "model_name": args.served_model_name, + "app_key": args.app_key, + } + self.kv_transfer_metrics = KVTransferMetrics(labels, args.metrics_reporters) + else: + self.kv_transfer_metrics = None + + # Build the Mooncake data-plane engine here and inject it into the + # transfer manager. self.kv_args is set above so register_buffer_to_engine + # (called by the base) sees the KV buffers. + engine = MooncakeTransferEngine( + hostname=get_local_ip_by_remote(), + gpu_id=kv_args.gpu_id, + ib_device=kv_args.ib_device, + ) + super().__init__(engine=engine) + + def register_buffer_to_engine(self): + for kv_data_ptr, kv_data_len in zip( + self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens + ): + self.engine.register(kv_data_ptr, kv_data_len) + for state_data_ptr, state_data_len in zip( + self.kv_args.state_data_ptrs, self.kv_args.state_data_lens + ): + self.engine.register(state_data_ptr, state_data_len) + + +class MooncakeKVBootstrapServer(DisaggBootstrapServerBase): + """KV bootstrap rendezvous: the shared server plus the prefill-side MLA / + kv-page length fields the decode side needs in the parallel-info sync.""" + + def __init__(self, port: int): + # Set before super() -- super() starts the server thread, after which a + # register PUT can call _ingest_put_extra and read these. + self.enable_mla_l1_5_cache = False + self.prefill_kv_item_lens = [] + self.prefill_kv_unit_lens = [] + self.prefill_state_item_lens = [] + self.prefill_state_unit_lens = [] + super().__init__(port) + + def _ingest_put_extra(self, data: dict) -> None: + self.enable_mla_l1_5_cache = bool(data["enable_mla_l1_5_cache"]) + self.prefill_kv_item_lens = data.get("kv_item_lens", self.prefill_kv_item_lens) + self.prefill_kv_unit_lens = data.get("kv_unit_lens", self.prefill_kv_unit_lens) + self.prefill_state_item_lens = data.get( + "state_item_lens", self.prefill_state_item_lens + ) + self.prefill_state_unit_lens = data.get( + "state_unit_lens", self.prefill_state_unit_lens + ) + + def _extra_parallel_info(self) -> dict: + return { + "enable_mla_l1_5_cache": self.enable_mla_l1_5_cache, + "kv_item_lens": self.prefill_kv_item_lens, + "kv_unit_lens": self.prefill_kv_unit_lens, + "state_item_lens": self.prefill_state_item_lens, + "state_unit_lens": self.prefill_state_unit_lens, + } diff --git a/python/tokenspeed/runtime/pd/mooncake/decode.py b/python/tokenspeed/runtime/pd/mooncake/decode.py new file mode 100644 index 0000000..013ee5f --- /dev/null +++ b/python/tokenspeed/runtime/pd/mooncake/decode.py @@ -0,0 +1,283 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import threading +import time +from collections import defaultdict +from dataclasses import dataclass + +import numpy as np +import requests + +from tokenspeed.runtime.pd.base.status import TransferPoll +from tokenspeed.runtime.pd.mooncake.conn import MooncakeKVManagerBase +from tokenspeed.runtime.pd.mooncake.entities import KVArgs, KVManagerArgs +from tokenspeed.runtime.pd.utils import DisaggregationMode +from tokenspeed.runtime.utils import ( + get_colorful_logger, +) +from tokenspeed.runtime.utils.env import envs +from tokenspeed.runtime.utils.network import get_free_port, get_local_ip_by_remote + +logger = get_colorful_logger(__name__) + + +@dataclass +class PrefillParallelInfo: + tp_size: int + dp_size: int + enable_mla_l1_5_cache: bool + kv_item_lens: tuple[int, ...] = () + kv_unit_lens: tuple[int, ...] = () + state_item_lens: tuple[int, ...] = () + state_unit_lens: tuple[int, ...] = () + + @property + def prefill_tp_size_per_dp_rank(self): + return self.tp_size // self.dp_size + + +def parse_prefill_status_message( + parts: list[bytes], +) -> tuple[int, int, int, int, list[int] | None]: + bootstrap_room = int(parts[0].decode("ascii")) + status = int(parts[1].decode("ascii")) + prefill_rank = int(parts[2].decode("ascii")) + bootstrap_token = int(parts[3].decode("ascii")) if len(parts) > 3 else -1 + spec_candidate_ids = None + if len(parts) > 4 and parts[4] != b"": + spec_candidate_ids = np.frombuffer(parts[4], dtype=np.int32).copy().tolist() + return ( + bootstrap_room, + status, + prefill_rank, + bootstrap_token, + spec_candidate_ids, + ) + + +class MooncakeKVManagerDecode(MooncakeKVManagerBase): + def __init__( + self, + args: KVManagerArgs, + kv_args: KVArgs, + ): + super().__init__(args, kv_args, 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() + # Heartbeat interval should be at least 2 seconds + self.heartbeat_interval = max( + envs.TOKENSPEED_DISAGGREGATION_HEARTBEAT_INTERVAL.get(), + 2.0, + ) + # Heartbeat failure should be at least 1 + self.max_failures = max( + envs.TOKENSPEED_DISAGGREGATION_HEARTBEAT_MAX_FAILURE.get(), 1 + ) + self.start_decode_thread() + self.connection_pool: dict[str, dict[str, str | int]] = {} + self.required_prefill_response_num_table: dict[int, int] = {} + self.prefill_response_tracker: dict[int, set[int]] = defaultdict(set) + + self.prefill_parallel_info: dict[str, PrefillParallelInfo] = {} + + # 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 = envs.TOKENSPEED_DISAGGREGATION_WAITING_TIMEOUT.get() + + def start_decode_thread(self): + self.rank_port = get_free_port() + self.server_socket.bind(f"tcp://{get_local_ip_by_remote()}:{self.rank_port}") + # Maps bootstrap_room -> bootstrap_token (first output token from prefill). + # Populated by decode_thread when a Success message carries a valid token, + # consumed by DisaggDecodeExecutor.generate_events() via pop_bootstrap_token(). + self.bootstrap_token_table: dict[int, int] = {} + self.spec_candidate_ids_table: dict[int, list[int]] = {} + self._pending_bootstrap_token_table: dict[int, int] = {} + self._pending_spec_candidate_ids_table: dict[int, list[int]] = {} + + def decode_thread(): + while True: + parts = self.server_socket.recv_multipart() + self._handle_prefill_status(*parse_prefill_status_message(parts)) + + def heartbeat_checker(): + while True: + time.sleep(self.heartbeat_interval) + with self.connection_lock: + addresses = list(self.prefill_parallel_info.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 TransferPoll.Success requests from the map + if bootstrap_room not in self.request_status: + self.addr_to_rooms_tracker[bootstrap_addr].discard( + bootstrap_room + ) + else: + logger.info( + "Attempting to reconnect to %s...", 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("Attempting to reconnect to %s...", 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 _handle_prefill_status( + self, + bootstrap_room: int, + status: int, + prefill_rank: int, + bootstrap_token: int, + spec_candidate_ids: list[int] | None, + ) -> None: + if status == TransferPoll.Success: + if bootstrap_room not in self.request_status: + return + + self.prefill_response_tracker[bootstrap_room].add(prefill_rank) + if bootstrap_token != -1: + self._pending_bootstrap_token_table.setdefault( + bootstrap_room, bootstrap_token + ) + if spec_candidate_ids is not None: + self._pending_spec_candidate_ids_table.setdefault( + bootstrap_room, spec_candidate_ids + ) + + expected_response_num = self.required_prefill_response_num_table.get( + bootstrap_room, 1 + ) + arrived_response_num = len(self.prefill_response_tracker[bootstrap_room]) + if arrived_response_num < expected_response_num: + return + + # Store metadata before marking Success so generate_events() can read + # it atomically in the same iteration. In heterogeneous TP, the rank + # carrying bootstrap metadata is not guaranteed to be the last rank + # whose transfer completes, so preserve the first valid value seen. + if bootstrap_room in self._pending_bootstrap_token_table: + self.bootstrap_token_table[bootstrap_room] = ( + self._pending_bootstrap_token_table.pop(bootstrap_room) + ) + if bootstrap_room in self._pending_spec_candidate_ids_table: + self.spec_candidate_ids_table[bootstrap_room] = ( + self._pending_spec_candidate_ids_table.pop(bootstrap_room) + ) + self.update_status(bootstrap_room, TransferPoll.Success) + return + + if status == TransferPoll.Failed: + self.record_failure( + bootstrap_room, + "Failed to get kvcache from prefill instance, it might be dead", + ) + self.update_status(bootstrap_room, TransferPoll.Failed) + return + + self.update_status(bootstrap_room, status) + + def pop_bootstrap_token(self, bootstrap_room: int) -> int: + """Pop and return the bootstrap_token for the given room, or -1 if absent.""" + return self.bootstrap_token_table.pop(bootstrap_room, -1) + + def pop_prefill_metadata(self, bootstrap_room: int) -> tuple[int, list[int] | None]: + return ( + self.bootstrap_token_table.pop(bootstrap_room, -1), + self.spec_candidate_ids_table.pop(bootstrap_room, None), + ) + + def get_session_id(self): + return self.engine.get_session_id() + + 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_parallel_info: + del self.prefill_parallel_info[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 TransferPoll.Failed + affected_rooms = [] + for room in possible_affected_rooms: + if ( + room in self.request_status + and self.check_status(room) != TransferPoll.Success + ): + self.record_failure( + room, + f"Losing connection with prefill instance (bootstrap_addr: {failed_bootstrap_addr})", + ) + self.update_status(room, TransferPoll.Failed) + affected_rooms.append(room) + logger.error( + "Losing connection with prefill instance (bootstrap_addr: %s), affected %s requests", + failed_bootstrap_addr, + len(affected_rooms), + ) diff --git a/python/tokenspeed/runtime/pd/mooncake/entities.py b/python/tokenspeed/runtime/pd/mooncake/entities.py new file mode 100644 index 0000000..d34f5eb --- /dev/null +++ b/python/tokenspeed/runtime/pd/mooncake/entities.py @@ -0,0 +1,237 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import dataclasses +import struct + +import numpy as np +import numpy.typing as npt + +from tokenspeed.runtime.pd.transfer_plan import ( + TransferFragment, + decode_transfer_fragments, +) +from tokenspeed.runtime.pd.utils import PageTransferMetadata + + +@dataclasses.dataclass +class KVArgs: + engine_rank: int + kv_data_ptrs: list[int] + kv_data_lens: list[int] + kv_item_lens: list[int] + offsets: list[tuple[int]] + aux_data_ptrs: list[int] + aux_data_lens: list[int] + aux_item_lens: list[int] + ib_device: str + gpu_id: int + target_layer_num: int + draft_layer_num: int + kv_layer_ids: list[int] = dataclasses.field(default_factory=list) + kv_unit_lens: list[int] = dataclasses.field(default_factory=list) + state_data_ptrs: list[int] = dataclasses.field(default_factory=list) + state_data_lens: list[int] = dataclasses.field(default_factory=list) + state_item_lens: list[int] = dataclasses.field(default_factory=list) + state_unit_lens: list[int] = dataclasses.field(default_factory=list) + state_type: str = "none" + state_layer_ids: list[int] = dataclasses.field(default_factory=list) + mamba_offsets: list[int] | None = None + + +class KVTransferError(Exception): + def __init__( + self, bootstrap_room: int, failure_reason: str, remote_endpoint: str = None + ): + super().__init__(failure_reason) + self.bootstrap_room = bootstrap_room + self.failure_reason = failure_reason + self.remote_endpoint = remote_endpoint + + def __str__(self): + if self.remote_endpoint: + return f"KVTransferError(bootstrap_room={self.bootstrap_room}, remote_endpoint={self.remote_endpoint}): {self.failure_reason}" + else: + return f"KVTransferError(bootstrap_room={self.bootstrap_room}): {self.failure_reason}" + + +# prefill +@dataclasses.dataclass +class TransferKVChunk: + room: int + prefill_kv_indices: npt.NDArray[np.int64] + index_slice: slice + is_last: bool + prefill_aux_index: int | None + mla_l1_5_args: PageTransferMetadata | None + prefill_mamba_indices: npt.NDArray[np.int64] | None = None + # First token generated by prefill (delivered via ZMQ, not RDMA). -1 = unavailable. + bootstrap_token: int = -1 + begin_cache_step: int | None = None + layerwise_interval: int = 1 + wait_for_bootstrap_token: bool = False + spec_candidate_ids: list[int] | None = None + + +@dataclasses.dataclass +class TransferIndexResolution: + src_indices: npt.NDArray[np.int64] + dst_indices: npt.NDArray[np.int64] + + +# decode +@dataclasses.dataclass +class TransferInfo: + room: int + endpoint: str + dst_port: int + mooncake_session_id: str + dst_kv_indices: npt.NDArray[np.int64] + dst_aux_index: int + required_dst_info_num: int + decode_prefix_len: int + dst_indices_are_local: bool + dst_page_transfer_mask: npt.NDArray[np.bool_] | None + dst_page_local_indices: npt.NDArray[np.int64] | None + dst_page_indices_mapping: npt.NDArray[np.int64] | None + dst_mamba_indices: npt.NDArray[np.int64] | None + is_dummy: bool + transfer_fragments: tuple[TransferFragment, ...] = () + + @classmethod + def from_zmq(cls, msg: list[bytes]): + transfer_fragments = () + if msg[4] == b"" and msg[5] == b"": + dst_kv_indices = np.array([], dtype=np.int64) + dst_aux_index = None + decode_prefix_len = 0 + dst_indices_are_local = False + dst_page_transfer_mask = None + dst_page_local_indices = None + dst_page_indices_mapping = None + dst_mamba_indices = None + is_dummy = True + else: + dst_kv_indices = np.frombuffer(msg[4], dtype=np.int64) + dst_aux_index = int(msg[5].decode("ascii")) + # decode_prefix_len is now msg[7] (we added it as additional message part) + decode_prefix_len = int(msg[7].decode("ascii")) if len(msg) > 7 else 0 + dst_indices_are_local = ( + bool(int(msg[8].decode("ascii"))) if len(msg) > 8 else False + ) + dst_page_transfer_mask = ( + np.frombuffer(msg[9], dtype=np.bool_) if len(msg) > 9 else None + ) + dst_page_local_indices = ( + np.frombuffer(msg[10], dtype=np.int64) if len(msg) > 10 else None + ) + dst_page_indices_mapping = ( + np.cumsum(dst_page_transfer_mask) - 1 + if dst_page_transfer_mask is not None + else None + ) + dst_mamba_indices = ( + np.frombuffer(msg[11], dtype=np.int64) + if len(msg) > 11 and msg[11] != b"" + else None + ) + transfer_fragments = ( + decode_transfer_fragments(msg[12], msg[13]) if len(msg) > 13 else () + ) + 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")), + decode_prefix_len=decode_prefix_len, + dst_indices_are_local=dst_indices_are_local, + dst_page_transfer_mask=dst_page_transfer_mask, + dst_page_local_indices=dst_page_local_indices, + dst_page_indices_mapping=dst_page_indices_mapping, + dst_mamba_indices=dst_mamba_indices, + is_dummy=is_dummy, + transfer_fragments=transfer_fragments, + ) + + +# 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_state_data_ptrs: list[int] + decode_prefix_len: int + + @classmethod + def from_zmq(cls, msg: list[bytes]): + # Format: room, endpoint, port, session, kv_ptrs, aux_ptrs, state_ptrs, decode_prefix_len. + # Older senders used msg[6] for decode_prefix_len and omitted state_ptrs. + has_state_frame = len(msg) >= 8 + if has_state_frame: + state_frame = msg[6] + decode_prefix_len = int(msg[7].decode("ascii")) if msg[7] else 0 + else: + state_frame = b"" + decode_prefix_len = ( + int(msg[6].decode("ascii")) if len(msg) >= 7 and msg[6] else 0 + ) + + 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_state_data_ptrs=list( + struct.unpack(f"{len(state_frame) // 8}Q", state_frame) + ), + decode_prefix_len=decode_prefix_len, + ) + + +@dataclasses.dataclass +class KVManagerArgs: + bootstrap_port: int + dist_init_addr: str + + world_size: int + dp_size: int + attn_tp_rank: int + attn_dp_rank: int + + is_mla_backend: bool + draft_is_mla_backend: bool + + enable_metrics: bool + enable_dp_attention: bool + enable_mla_l1_5_cache: bool + + served_model_name: str + app_key: str + metrics_reporters: str diff --git a/python/tokenspeed/runtime/pd/mooncake/prefill.py b/python/tokenspeed/runtime/pd/mooncake/prefill.py new file mode 100644 index 0000000..8dd480d --- /dev/null +++ b/python/tokenspeed/runtime/pd/mooncake/prefill.py @@ -0,0 +1,1100 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import concurrent.futures +import os +import socket +import threading +import time +from collections import defaultdict + +import numpy as np +import numpy.typing as npt +import requests + +from tokenspeed.runtime.pd.base.status import TransferPoll +from tokenspeed.runtime.pd.mooncake.conn import MooncakeKVManagerBase +from tokenspeed.runtime.pd.mooncake.entities import ( + KVArgs, + KVArgsRegisterInfo, + KVManagerArgs, + TransferIndexResolution, + TransferInfo, + TransferKVChunk, +) +from tokenspeed.runtime.pd.transfer_plan import ( + BufferKind, + TransferFragment, +) +from tokenspeed.runtime.pd.utils import ( + DisaggregationMode, + FastQueue, + PageTransferMetadata, + StepCounter, + group_concurrent_contiguous, +) +from tokenspeed.runtime.utils import ( + get_colorful_logger, +) +from tokenspeed.runtime.utils.env import envs +from tokenspeed.runtime.utils.network import ( + get_free_port, + get_ip, + get_local_ip_by_remote, +) + +logger = get_colorful_logger(__name__) + + +class MooncakeKVManagerPrefill(MooncakeKVManagerBase): + def __init__( + self, + args: KVManagerArgs, + kv_args: KVArgs, + ): + super().__init__(args, kv_args, 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: dict[str, float] = {} + self.failed_session_ttl = max( + envs.TOKENSPEED_DISAGGREGATION_FAILED_SESSION_TTL.get(), 0 + ) + self.session_lock = threading.Lock() + self.kv_layer_ids = list( + getattr(self.kv_args, "kv_layer_ids", None) + or range(len(self.kv_args.offsets)) + ) + self.state_layer_ids = list(getattr(self.kv_args, "state_layer_ids", []) or []) + layer_ids = self.kv_layer_ids + self.state_layer_ids + self.layer_num = ( + (max(layer_ids) + 1) if layer_ids else len(self.kv_args.offsets) + ) + self._kv_layer_to_index = { + layer_id: i + for i, layer_id in enumerate(self.kv_layer_ids[: len(self.kv_args.offsets)]) + } + self.layerwise_interval = 1 + self.layerwise_debug = envs.TOKENSPEED_PD_LAYERWISE_DEBUG.get() + self.step_counter = None + # room -> (bootstrap_token, spec_candidate_ids). Published after the prefill + # forward; the transfer thread reads it on the wait_for_bootstrap_token path. + self.prefill_metadata: dict[int, tuple[int, list[int] | None]] = {} + self.expired_prefill_metadata_rooms: set[int] = set() + self.bootstrap_token_cond = threading.Condition() + # Determine the number of threads to use for kv sender + cpu_count = os.cpu_count() + transfer_thread_pool_size = ( + envs.TOKENSPEED_DISAGGREGATION_THREAD_POOL_SIZE.get_set_value_or( + min(max(4, int(0.75 * cpu_count) // 8), 12) + ) + ) + transfer_queue_size = envs.TOKENSPEED_DISAGGREGATION_QUEUE_SIZE.get() + if transfer_thread_pool_size < transfer_queue_size: + raise ValueError( + "TOKENSPEED_DISAGGREGATION_THREAD_POOL_SIZE=" + f"{transfer_thread_pool_size} must be greater than or equal to " + f"TOKENSPEED_DISAGGREGATION_QUEUE_SIZE={transfer_queue_size}." + ) + self.start_transfer_thread(transfer_thread_pool_size, transfer_queue_size) + self.bootstrap_time_out = envs.TOKENSPEED_DISAGGREGATION_BOOTSTRAP_TIMEOUT.get() + + def register_layerwise_step_counter( + self, step_counter: StepCounter, interval: int + ) -> None: + self.step_counter = step_counter + self.layerwise_interval = max(int(interval), 1) + + def reserve_layerwise_cache_steps(self) -> int: + if self.step_counter is None: + return 0 + cache_step, _ = self.step_counter.current_step() + self.step_counter.advance_step( + delta_cache_step=self.layer_num, + delta_aux_step=0, + ) + return cache_step + + def set_prefill_metadata( + self, + room: int, + token: int, + spec_candidate_ids: list[int] | None = None, + ) -> None: + with self.bootstrap_token_cond: + if room in self.expired_prefill_metadata_rooms: + self.expired_prefill_metadata_rooms.discard(room) + logger.warning( + "Dropping late prefill metadata for expired bootstrap_room=%s", + room, + ) + return + self.prefill_metadata[room] = ( + token, + spec_candidate_ids, + ) + self.bootstrap_token_cond.notify_all() + + def discard_expired_metadata_room(self, room: int) -> None: + """Best-effort cleanup of per-room metadata and expiry markers.""" + with self.bootstrap_token_cond: + self.expired_prefill_metadata_rooms.discard(room) + self.prefill_metadata.pop(room, None) + + def _wait_prefill_metadata( + self, + room: int | None, + fallback_token: int, + fallback_candidate_ids: list[int] | None, + ) -> tuple[int, list[int] | None]: + if room is None or fallback_token != -1: + return fallback_token, fallback_candidate_ids + wait_log_interval = max(envs.TOKENSPEED_PD_PREFILL_METADATA_TIMEOUT.get(), 0.01) + start_time = time.monotonic() + next_log_time = start_time + wait_log_interval + with self.bootstrap_token_cond: + while room not in self.prefill_metadata: + if self.request_status.get(room) == TransferPoll.Failed: + self.expired_prefill_metadata_rooms.add(room) + logger.warning( + "Prefill metadata unavailable for failed bootstrap_room=%s; using fallback=%s", + room, + fallback_token, + ) + return ( + fallback_token, + fallback_candidate_ids, + ) + now = time.monotonic() + if now >= next_log_time: + logger.debug( + "Still waiting for prefill metadata for bootstrap_room=%s after %.2fs", + room, + now - start_time, + ) + next_log_time = now + wait_log_interval + self.bootstrap_token_cond.wait(timeout=0.01) + return self.prefill_metadata[room] + + def _is_session_failed(self, mooncake_session_id: str) -> bool: + if self.failed_session_ttl <= 0: + return False + failed_at = self.failed_sessions.get(mooncake_session_id) + if failed_at is None: + return False + elapsed = time.monotonic() - failed_at + logger.info( + "Session %s failed for %.2fs (TTL=%ds).", + mooncake_session_id, + elapsed, + self.failed_session_ttl, + ) + if elapsed < self.failed_session_ttl: + return True + del self.failed_sessions[mooncake_session_id] + logger.info( + "Session %s failed TTL expired (%.2fs >= %ds), reset.", + mooncake_session_id, + elapsed, + self.failed_session_ttl, + ) + return False + + def _mark_session_failed( + self, mooncake_session_id: str, reason: str = "transfer_failed" + ) -> None: + if self.failed_session_ttl <= 0: + return + self.failed_sessions[mooncake_session_id] = time.monotonic() + logger.warning( + "Session %s marked failed (reason=%s, ttl=%ds).", + mooncake_session_id, + reason, + self.failed_session_ttl, + ) + + def _clear_failed_session(self, mooncake_session_id: str) -> None: + if mooncake_session_id in self.failed_sessions: + del self.failed_sessions[mooncake_session_id] + logger.info( + "Session %s failed state cleared due to KVArgs registration.", + mooncake_session_id, + ) + if mooncake_session_id in self.session_failures: + del self.session_failures[mooncake_session_id] + + def resolve_transfer_indices( + self, + kv_chunk: TransferKVChunk, + req: TransferInfo, + ) -> TransferIndexResolution: + src_indices = kv_chunk.prefill_kv_indices + dst_indices = req.dst_kv_indices[kv_chunk.index_slice] + + valid_len = min(len(src_indices), len(dst_indices)) + # Fast path: empty transfer chunk. Avoid MLA assertions/index ops on empty payload. + if valid_len == 0: + empty = np.array([], dtype=np.int64) + return TransferIndexResolution(src_indices=empty, dst_indices=empty) + + if valid_len < len(src_indices) or valid_len < len(dst_indices): + logger.warning( + "Mismatched transfer indices, truncating to %s (src=%s, dst=%s)", + valid_len, + len(src_indices), + len(dst_indices), + ) + src_indices = src_indices[:valid_len] + dst_indices = dst_indices[:valid_len] + + src_mode = self.src_mode + dst_mode = "ON" if req.dst_indices_are_local else "OFF" + src_args = kv_chunk.mla_l1_5_args + + # Prefill OFF/Decode OFF: use original indices. + if src_mode == "OFF" and dst_mode == "OFF": + return TransferIndexResolution(src_indices, dst_indices) + + # Prefill ON/Decode OFF: prefill-side only has partial kv cache, only send local part. + if src_mode == "ON" and dst_mode == "OFF": + if src_args is None: + raise RuntimeError( + "Prefill MLA L1.5 cache is enabled but no transfer " + "metadata provided" + ) + src_mask = src_args.page_transfer_mask + src_local = src_args.page_local_indices + return TransferIndexResolution( + src_indices=src_local, + dst_indices=dst_indices[src_mask], + ) + + # Prefill OFF/Decode ON: decode-side only has partial kv cache, only send requested part. + if src_mode == "OFF" and dst_mode == "ON": + if req.dst_page_transfer_mask is None: + raise RuntimeError( + "OFF/ON expects decode ownership mask when destination " + "reports local indices." + ) + dst_mapping = req.dst_page_indices_mapping[kv_chunk.index_slice] + dst_mask = req.dst_page_transfer_mask[kv_chunk.index_slice] + dst_mapping = dst_mapping[dst_mask] + dst_local = req.dst_page_local_indices[dst_mapping] + return TransferIndexResolution( + src_indices=src_indices[dst_mask], + dst_indices=dst_local, + ) + + # Prefill ON/Decode ON: both sides hold partial kv cache, find the intersection part. + if src_args is None: + raise RuntimeError( + "ON/ON expects prefill metadata " + "(page_transfer_mask/page_local_indices)." + ) + if req.dst_page_transfer_mask is None: + raise RuntimeError( + "ON/ON expects decode ownership mask when destination uses local " + "index space." + ) + + # src_args.page_transfer_mask is generated from current chunk page_indices, + # so it is already in chunk-local coordinates. + src_mask = src_args.page_transfer_mask[:valid_len] + src_local = src_args.page_local_indices + + # req.dst_page_transfer_mask is request-global and should be sliced by chunk. + dst_mask = req.dst_page_transfer_mask[kv_chunk.index_slice][:valid_len] + + # Only positions owned by both sides should be transferred. + common_mask = src_mask & dst_mask + + dst_mapping = req.dst_page_indices_mapping[kv_chunk.index_slice][:valid_len] + dst_mapping = dst_mapping[common_mask] + dst_local = req.dst_page_local_indices[dst_mapping] + + src_mapping = np.cumsum(src_args.page_transfer_mask) - 1 + src_mapping = src_mapping[:valid_len] + src_mapping = src_mapping[common_mask] + src_local = src_args.page_local_indices[src_mapping] + + return TransferIndexResolution( + src_indices=src_local, + dst_indices=dst_local, + ) + + 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.int64], + dst_kv_ptrs: list[int], + dst_kv_indices: npt.NDArray[np.int64], + executor: concurrent.futures.ThreadPoolExecutor, + transfer_fragments: tuple[TransferFragment, ...] = (), + ): + # Group by indices + prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous( + prefill_kv_indices, dst_kv_indices + ) + + # ``_layer_transfer_blocks`` indexes ``kv_args.offsets`` by its loop var, + # and ``offsets`` is keyed by KV-LAYER INDEX (one entry per attention/KV + # layer), not by global layer id. For hybrid models (e.g. Qwen3.5 GDN + + # attention) only the attention layers carry KV, so len(offsets) is the + # KV-layer count while ``self.layer_num`` is the (larger) global layer + # count -- using ``layer_num`` here over-runs ``offsets`` and IndexErrors. + # Iterate the full KV-index space instead (the layerwise path already + # maps global->KV index via ``_kv_layer_to_index`` before calling this). + transfer_blocks = self._layer_transfer_blocks( + dst_ptrs=dst_kv_ptrs, + src_blocks=prefill_kv_blocks, + dst_blocks=dst_kv_blocks, + begin_layer_id=0, + end_layer_id=len(self.kv_args.offsets), + transfer_fragments=transfer_fragments, + ) + return self._transfer_data(mooncake_session_id, transfer_blocks) + + def _layer_transfer_blocks( + self, + dst_ptrs: list[int], + src_blocks, + dst_blocks, + begin_layer_id: int, + end_layer_id: int, + transfer_fragments: tuple[TransferFragment, ...] = (), + ) -> list[tuple[int, int, int]]: + transfer_blocks = [] + fragments_by_buffer: dict[int, list[TransferFragment]] = defaultdict(list) + for fragment in transfer_fragments: + if fragment.buffer_kind != BufferKind.MAMBA_STATE: + fragments_by_buffer[fragment.buffer_index].append(fragment) + has_fragment_plan = bool(transfer_fragments) + + for layer_id in range(begin_layer_id, end_layer_id): + for ptr_offset in self.kv_args.offsets[layer_id]: + src_ptr = self.kv_args.kv_data_ptrs[ptr_offset] + dst_ptr = dst_ptrs[ptr_offset] + item_len = self.kv_args.kv_item_lens[ptr_offset] + buffer_fragments = fragments_by_buffer.get(ptr_offset, ()) + if has_fragment_plan: + for fragment in buffer_fragments: + for prefill_index, decode_index in zip(src_blocks, dst_blocks): + page_count = min(len(prefill_index), len(decode_index)) + if fragment.page_count is not None: + page_count = min(page_count, fragment.page_count) + for page_offset in range(page_count): + src_addr = ( + src_ptr + + int(prefill_index[page_offset]) + * fragment.src_page_stride_bytes + + fragment.src_byte_offset + ) + dst_addr = ( + dst_ptr + + int(decode_index[page_offset]) + * fragment.dst_page_stride_bytes + + fragment.dst_byte_offset + ) + transfer_blocks.append( + (src_addr, dst_addr, fragment.bytes_per_page) + ) + continue + + for prefill_index, decode_index in zip(src_blocks, dst_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 + + def _wait_until_cache_step(self, target_step: int) -> None: + if self.step_counter is None: + return + while True: + ready_step = self.step_counter.query_ready_cache_step() + if StepCounter.is_step_ready(ready_step, target_step): + return + time.sleep(1e-4) + + def send_mamba_cache( + self, + mooncake_session_id: str, + prefill_mamba_indices: npt.NDArray[np.int64] | None, + dst_state_data_ptrs: list[int], + dst_mamba_indices: npt.NDArray[np.int64] | None, + begin_layer_id: int | None = None, + end_layer_id: int | None = None, + transfer_fragments: tuple[TransferFragment, ...] = (), + ) -> int: + if self.kv_args.state_type != "mamba": + return 0 + state_ptrs = self.kv_args.state_data_ptrs + state_item_lens = self.kv_args.state_item_lens + if ( + not state_ptrs + or not dst_state_data_ptrs + or prefill_mamba_indices is None + or dst_mamba_indices is None + ): + return 0 + if len(state_ptrs) != len(dst_state_data_ptrs): + logger.error( + "Mamba state tensor count mismatch: prefill=%d decode=%d", + len(state_ptrs), + len(dst_state_data_ptrs), + ) + return -1 + + if prefill_mamba_indices.shape != dst_mamba_indices.shape: + if prefill_mamba_indices.size == 1 and dst_mamba_indices.size > 1: + prefill_mamba_indices = np.full( + dst_mamba_indices.shape, + int(prefill_mamba_indices[0]), + dtype=np.int64, + ) + else: + logger.error( + "Mamba state slot count mismatch: prefill=%s decode=%s", + prefill_mamba_indices.tolist(), + dst_mamba_indices.tolist(), + ) + return -1 + + state_items = list( + enumerate(zip(state_ptrs, dst_state_data_ptrs, state_item_lens)) + ) + if begin_layer_id is not None or end_layer_id is not None: + begin = 0 if begin_layer_id is None else begin_layer_id + end = self.layer_num if end_layer_id is None else end_layer_id + if len(self.state_layer_ids) != len(state_items): + logger.error( + "Mamba state layer id count mismatch: ids=%d tensors=%d", + len(self.state_layer_ids), + len(state_items), + ) + return -1 + state_items = [ + item + for item, layer_id in zip(state_items, self.state_layer_ids) + if begin <= layer_id < end + ] + if not state_items: + return 0 + + valid = (prefill_mamba_indices >= 0) & (dst_mamba_indices >= 0) + log_layerwise = getattr(self, "layerwise_debug", False) + if log_layerwise and begin_layer_id is not None and end_layer_id is not None: + logger.info( + "[layerwise_transfer] session=%s layers=[%d,%d) " + "send mamba tensors=%d bytes=%d", + mooncake_session_id, + begin_layer_id, + end_layer_id, + len(state_items), + sum(item_len for _, (_, _, item_len) in state_items) * int(valid.sum()), + ) + if not valid.any(): + return 0 + + src_indices = prefill_mamba_indices[valid] + dst_indices = dst_mamba_indices[valid] + src_blocks, dst_blocks = group_concurrent_contiguous(src_indices, dst_indices) + transfer_blocks = [] + state_fragments_by_buffer: dict[int, list[TransferFragment]] = defaultdict(list) + for fragment in transfer_fragments: + if fragment.buffer_kind == BufferKind.MAMBA_STATE: + state_fragments_by_buffer[fragment.buffer_index].append(fragment) + has_state_fragment_plan = bool(state_fragments_by_buffer) + + for state_index, (src_ptr, dst_ptr, item_len) in state_items: + buffer_fragments = state_fragments_by_buffer.get(state_index, ()) + if has_state_fragment_plan: + for fragment in buffer_fragments: + for prefill_index, decode_index in zip(src_blocks, dst_blocks): + page_count = min(len(prefill_index), len(decode_index)) + if fragment.page_count is not None: + page_count = min(page_count, fragment.page_count) + for page_offset in range(page_count): + src_addr = ( + src_ptr + + int(prefill_index[page_offset]) + * fragment.src_page_stride_bytes + + fragment.src_byte_offset + ) + dst_addr = ( + dst_ptr + + int(decode_index[page_offset]) + * fragment.dst_page_stride_bytes + + fragment.dst_byte_offset + ) + transfer_blocks.append( + (src_addr, dst_addr, fragment.bytes_per_page) + ) + continue + + for prefill_index, decode_index in zip(src_blocks, dst_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)) + + total_bytes = sum(length for _, _, length in transfer_blocks) + ret = self._transfer_data(mooncake_session_id, transfer_blocks) + logger.debug( + "Transferred mamba cache for session=%s slots=%s blocks=%d bytes=%d ret=%s", + mooncake_session_id, + src_indices.tolist(), + len(transfer_blocks), + total_bytes, + ret, + ) + return ret + + def send_kvcache_layerwise( + self, + mooncake_session_id: str, + prefill_kv_indices: npt.NDArray[np.int64], + dst_kv_ptrs: list[int], + dst_kv_indices: npt.NDArray[np.int64], + begin_cache_step: int, + interval: int, + dst_state_data_ptrs: list[int] | None = None, + prefill_mamba_indices: npt.NDArray[np.int64] | None = None, + dst_mamba_indices: npt.NDArray[np.int64] | None = None, + transfer_fragments: tuple[TransferFragment, ...] = (), + ) -> int: + prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous( + prefill_kv_indices, dst_kv_indices + ) + + interval = max(int(interval), 1) + log_layerwise = getattr(self, "layerwise_debug", False) + for begin_layer_id in range(0, self.layer_num, interval): + end_layer_id = min(begin_layer_id + interval, self.layer_num) + target_step = begin_cache_step + end_layer_id - 1 + if log_layerwise: + logger.info( + "[layerwise_transfer] session=%s layers=[%d,%d) wait_cache_step=%d pages=%d", + mooncake_session_id, + begin_layer_id, + end_layer_id, + target_step, + len(prefill_kv_indices), + ) + self._wait_until_cache_step(target_step) + + transfer_blocks = [] + if prefill_kv_blocks: + for global_layer_id in range(begin_layer_id, end_layer_id): + kv_layer_index = self._kv_layer_to_index.get(global_layer_id) + if kv_layer_index is None: + continue + transfer_blocks.extend( + self._layer_transfer_blocks( + dst_ptrs=dst_kv_ptrs, + src_blocks=prefill_kv_blocks, + dst_blocks=dst_kv_blocks, + begin_layer_id=kv_layer_index, + end_layer_id=kv_layer_index + 1, + transfer_fragments=transfer_fragments, + ) + ) + if transfer_blocks: + if log_layerwise: + total_bytes = sum(length for _, _, length in transfer_blocks) + logger.info( + "[layerwise_transfer] session=%s layers=[%d,%d) send kv blocks=%d bytes=%d", + mooncake_session_id, + begin_layer_id, + end_layer_id, + len(transfer_blocks), + total_bytes, + ) + ret = self._transfer_data(mooncake_session_id, transfer_blocks) + if ret != 0: + return ret + + ret = self.send_mamba_cache( + mooncake_session_id, + prefill_mamba_indices, + dst_state_data_ptrs or [], + dst_mamba_indices, + begin_layer_id=begin_layer_id, + end_layer_id=end_layer_id, + transfer_fragments=transfer_fragments, + ) + if ret != 0: + return ret + if log_layerwise: + logger.info( + "[layerwise_transfer] session=%s layers=[%d,%d) done", + mooncake_session_id, + begin_layer_id, + end_layer_id, + ) + return 0 + + def sync_status_to_decode_endpoint( + self, + remote: str, + dst_port: int, + room: int, + status: int, + prefill_rank: int, + bootstrap_token: int = -1, + spec_candidate_ids: list[int] | None = None, + ): + if ":" in remote: + remote = remote.split(":")[0] + spec_candidate_payload = ( + np.asarray(spec_candidate_ids, dtype=np.int32).tobytes() + if spec_candidate_ids is not None + else b"" + ) + self._connect("tcp://" + remote + ":" + str(dst_port)).send_multipart( + [ + str(room).encode("ascii"), + str(status).encode("ascii"), + str(prefill_rank).encode("ascii"), + str(bootstrap_token).encode("ascii"), + spec_candidate_payload, + ] + ) + + def abort_room(self, room: int, reason: str) -> None: + """Notify the decode that a room failed before any KV transfer. + + EPD: when the prefill aborts a request on embedding-receive timeout it never + sends KV, so the decode's dual-dispatched KV receiver would wait + indefinitely -- its heartbeat only trips if the prefill /health dies, and the + receiver waiting_timeout only covers the WaitingForInput state (a receiver + whose prefill never registered a sender is stuck earlier). Push a Failed + status to every decode endpoint that already pre-allocated for this room + (mirrors the in-transfer failure path), so the decode raises a FailedEvent and + the client gets an error instead of hanging. A room whose decode has not + pre-allocated yet is only marked Failed locally (no endpoint to notify). + """ + self.record_failure(room, reason) + self.update_status(room, TransferPoll.Failed) + for req in list(self.transfer_infos.get(room, {}).values()): + if not req.is_dummy: + self.sync_status_to_decode_endpoint( + req.endpoint, + req.dst_port, + req.room, + TransferPoll.Failed, + self.attn_tp_rank, + ) + + def transfer_worker( + self, queue: FastQueue, executor: concurrent.futures.ThreadPoolExecutor + ): + while True: + try: + kv_chunk = queue.get() + logger.debug( + "[TRANSFER_WORKER] Got transfer request for room %s, is_last=%s, kv_indices_len=%s", + kv_chunk.room, + kv_chunk.is_last, + len(kv_chunk.prefill_kv_indices), + ) + reqs_to_be_processed = ( + self.transfer_infos[kv_chunk.room].values() + if kv_chunk.room in self.transfer_infos + else [] + ) + polls = [] + dst_ranks_infos = [] + for req in reqs_to_be_processed: + if not req.is_dummy: + # Early exit if the request has failed + with self.session_lock: + if self._is_session_failed(req.mooncake_session_id): + logger.info( + "Blocked transfer due to failed session (room=%s, session=%s).", + kv_chunk.room, + req.mooncake_session_id, + ) + 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, TransferPoll.Failed) + self.sync_status_to_decode_endpoint( + req.endpoint, + req.dst_port, + req.room, + TransferPoll.Failed, + self.attn_tp_rank, + ) + break + resolved = self.resolve_transfer_indices(kv_chunk, req) + + logger.debug( + "[TRANSFER_WORKER] Calling send_kvcache for room %s, session %s", + kv_chunk.room, + req.mooncake_session_id, + ) + tm_start = time.monotonic() + prefill_metadata = None + dst_kv_ptrs = self.decode_kv_args_table[ + req.mooncake_session_id + ].dst_kv_ptrs + dst_state_data_ptrs = self.decode_kv_args_table[ + req.mooncake_session_id + ].dst_state_data_ptrs + if kv_chunk.begin_cache_step is None: + ret = self.send_kvcache( + req.mooncake_session_id, + resolved.src_indices, + dst_kv_ptrs, + resolved.dst_indices, + executor, + req.transfer_fragments, + ) + else: + ret = self.send_kvcache_layerwise( + req.mooncake_session_id, + resolved.src_indices, + dst_kv_ptrs, + resolved.dst_indices, + kv_chunk.begin_cache_step, + kv_chunk.layerwise_interval, + dst_state_data_ptrs, + kv_chunk.prefill_mamba_indices, + req.dst_mamba_indices, + req.transfer_fragments, + ) + if ret == 0 and kv_chunk.is_last: + if kv_chunk.wait_for_bootstrap_token: + # The first decode/target-verify step consumes prefill's + # sampled bootstrap token and spec candidates; publish + # Success only after that metadata is ready. + prefill_metadata = self._wait_prefill_metadata( + kv_chunk.room, + kv_chunk.bootstrap_token, + kv_chunk.spec_candidate_ids, + ) + if kv_chunk.begin_cache_step is None: + ret = self.send_mamba_cache( + req.mooncake_session_id, + kv_chunk.prefill_mamba_indices, + dst_state_data_ptrs, + req.dst_mamba_indices, + transfer_fragments=req.transfer_fragments, + ) + logger.debug( + "[TRANSFER_WORKER] send_kvcache returned %s for room %s", + ret, + kv_chunk.room, + ) + 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._mark_session_failed( + req.mooncake_session_id, reason="send_kvcache" + ) + logger.error( + "Session %s failed.", req.mooncake_session_id + ) + 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, TransferPoll.Failed) + self.sync_status_to_decode_endpoint( + req.endpoint, + req.dst_port, + req.room, + TransferPoll.Failed, + self.attn_tp_rank, + ) + break + + if kv_chunk.is_last: + polls.append(True) + 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 = ( + TransferPoll.Success + if all(polls) + else TransferPoll.Failed + ) + self.update_status(req.room, status) + # bootstrap_token is carried directly in the chunk (set by + # DisaggPrefillExecutor._decode after prefill forward). + if kv_chunk.wait_for_bootstrap_token: + if prefill_metadata is None: + prefill_metadata = self._wait_prefill_metadata( + kv_chunk.room, + kv_chunk.bootstrap_token, + kv_chunk.spec_candidate_ids, + ) + bootstrap_token, spec_candidate_ids = ( + prefill_metadata + ) + else: + bootstrap_token, spec_candidate_ids = ( + kv_chunk.bootstrap_token, + kv_chunk.spec_candidate_ids, + ) + if self.check_status(req.room) == TransferPoll.Failed: + status = TransferPoll.Failed + for endpoint, dst_port, room in dst_ranks_infos: + self.sync_status_to_decode_endpoint( + endpoint, + dst_port, + room, + status, + self.attn_tp_rank, + bootstrap_token=bootstrap_token, + spec_candidate_ids=spec_candidate_ids, + ) + elapsed_seconds = time.monotonic() - tm_start + if self.kv_transfer_metrics: + self.kv_transfer_metrics.observe_kv_transfer_latency( + elapsed_seconds + ) + 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, TransferPoll.Success) + + if ( + kv_chunk.room not in self.request_status + or self.check_status(kv_chunk.room) == TransferPoll.Success + ): + if kv_chunk.room in self.transfer_infos: + self.transfer_infos.pop(kv_chunk.room) + + except Exception as exc: + raise RuntimeError( + f"Transfer thread failed because of {exc}. Prefill instance " + f"with bootstrap_port={self.bootstrap_port} is dead." + ) from exc + + def start_prefill_thread(self): + self.rank_port = get_free_port() + self.server_socket.bind(f"tcp://{get_local_ip_by_remote()}:{self.rank_port}") + + def bootstrap_thread(): + """This thread recvs pre-alloc notification from the decode engine""" + # TransferPoll.Bootstrapping -> TransferPoll.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") + logger.info( + "[Prefill bootstrap_thread] recv multipart: room=%s session_id=%s", + room, + mooncake_session_id, + ) + if room == "None": + self.decode_kv_args_table[mooncake_session_id] = ( + KVArgsRegisterInfo.from_zmq(waiting_req_bytes) + ) + with self.session_lock: + self._clear_failed_session(mooncake_session_id) + logger.info( + "[Prefill bootstrap_thread] registered kv_args from decode session=%s", + mooncake_session_id, + ) + 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) + ) + logger.info( + "[Prefill bootstrap_thread] pre-alloc received: room=%d session=%s got=%d/%d, status -> %s", + room, + mooncake_session_id, + len(self.transfer_infos[room]), + required_dst_info_num, + ( + "Bootstrapped" + if len(self.transfer_infos[room]) == required_dst_info_num + else "waiting more" + ), + ) + if len(self.transfer_infos[room]) == required_dst_info_num: + self.update_status(room, TransferPoll.Bootstrapped) + + threading.Thread(target=bootstrap_thread).start() + + def start_transfer_thread( + self, transfer_thread_pool_size: int, transfer_queue_size: int + ): + self.transfer_queues: list[FastQueue] = [ + FastQueue() for _ in range(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() + + def add_transfer_request( + self, + bootstrap_room: int, + kv_indices: npt.NDArray[np.int64], + index_slice: slice, + is_last: bool, + aux_index: int | None = None, + mla_l1_5_args: PageTransferMetadata | None = None, + bootstrap_token: int = -1, + begin_cache_step: int | None = None, + layerwise_interval: int = 1, + wait_for_bootstrap_token: bool = False, + mamba_indices: npt.NDArray[np.int64] | None = None, + spec_candidate_ids: list[int] | None = None, + ): + if self.disaggregation_mode != DisaggregationMode.PREFILL: + raise RuntimeError("Transfer requests can only be added in prefill mode.") + if is_last and aux_index is None: + raise ValueError("aux_index must be set for the last transfer chunk.") + if ( + bootstrap_room not in self.request_status + or self.check_status(bootstrap_room) == TransferPoll.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 + + # 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.split(":")[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, + mla_l1_5_args=mla_l1_5_args, + bootstrap_token=bootstrap_token, + begin_cache_step=begin_cache_step, + layerwise_interval=layerwise_interval, + wait_for_bootstrap_token=wait_for_bootstrap_token, + prefill_mamba_indices=mamba_indices, + spec_candidate_ids=spec_candidate_ids, + ) + ) + + def receive_decode_prefix_info(self, bootstrap_room: int) -> int: + """Receive decode prefix info from decode side""" + # In mooncake implementation, decode_prefix_len is handled via ZMQ messages + # Check the stored transfer info for this room + if bootstrap_room in self.transfer_infos: + for transfer_info in self.transfer_infos[bootstrap_room].values(): + if ( + hasattr(transfer_info, "decode_prefix_len") + and transfer_info.decode_prefix_len > 0 + ): + logger.debug( + "Found decode_prefix_len=%s for room %s", + transfer_info.decode_prefix_len, + bootstrap_room, + ) + return transfer_info.decode_prefix_len + logger.debug("No decode_prefix_len found for room %s, using 0", bootstrap_room) + return 0 + + def _register_to_bootstrap(self): + """Register KVSender to bootstrap server via HTTP POST.""" + if self.dist_init_addr: + ip_address = socket.gethostbyname(self.dist_init_addr.split(":")[0]) + else: + ip_address = get_ip() + + bootstrap_server_url = f"{ip_address}:{self.bootstrap_port}" + url = f"http://{bootstrap_server_url}/route" + payload = { + "role": "Prefill", + "world_size": self.world_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, + "enable_mla_l1_5_cache": self.args.enable_mla_l1_5_cache, + "kv_item_lens": self.kv_args.kv_item_lens, + "kv_unit_lens": getattr(self.kv_args, "kv_unit_lens", []), + "state_item_lens": self.kv_args.state_item_lens, + "state_unit_lens": getattr(self.kv_args, "state_unit_lens", []), + } + + 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( + "Prefill instance failed to connect to bootstrap server: %s, %s", + response.status_code, + response.text, + ) + except Exception as exc: + logger.error( + "Prefill instance failed to register with bootstrap server: %s", exc + ) + + +from tokenspeed.runtime.pd.mooncake.sender import MooncakeKVSender + +__all__ = ["MooncakeKVManagerPrefill", "MooncakeKVSender"] diff --git a/python/tokenspeed/runtime/pd/mooncake/receiver.py b/python/tokenspeed/runtime/pd/mooncake/receiver.py new file mode 100644 index 0000000..d9472ee --- /dev/null +++ b/python/tokenspeed/runtime/pd/mooncake/receiver.py @@ -0,0 +1,730 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import struct +import threading +import time +from dataclasses import dataclass + +import numpy as np +import numpy.typing as npt +import requests +import zmq + +from tokenspeed.runtime.pd.base.status import TransferPoll +from tokenspeed.runtime.pd.mooncake.entities import KVTransferError +from tokenspeed.runtime.pd.transfer_plan import ( + BufferKind, + BufferLayout, + ParallelLayout, + PDTransferPlanner, + RankTransferPlan, + encode_transfer_fragments, +) +from tokenspeed.runtime.pd.utils import ( + PageTransferMetadata, +) +from tokenspeed.runtime.utils import ( + get_colorful_logger, +) +from tokenspeed.runtime.utils.network import get_local_ip_by_remote + +logger = get_colorful_logger(__name__) + +from tokenspeed.runtime.pd.mooncake.decode import ( + MooncakeKVManagerDecode, + PrefillParallelInfo, +) + + +def _get_prefill_parallel_info_from_server( + bootstrap_addr, +) -> PrefillParallelInfo | None: + """Fetch the prefill parallel info from the bootstrap server.""" + try: + url = f"http://{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 PrefillParallelInfo( + tp_size=int(prefill_parallel_info["prefill_tp_size"]), + dp_size=int(prefill_parallel_info["prefill_dp_size"]), + enable_mla_l1_5_cache=bool( + prefill_parallel_info["enable_mla_l1_5_cache"] + ), + kv_item_lens=tuple( + int(x) for x in prefill_parallel_info.get("kv_item_lens", []) + ), + kv_unit_lens=tuple( + int(x) for x in prefill_parallel_info.get("kv_unit_lens", []) + ), + state_item_lens=tuple( + int(x) for x in prefill_parallel_info.get("state_item_lens", []) + ), + state_unit_lens=tuple( + int(x) for x in prefill_parallel_info.get("state_unit_lens", []) + ), + ) + else: + logger.error( + "Failed to get prefill parallel info: %s, %s", + response.status_code, + response.text, + ) + return None + except Exception as exc: + logger.error("Error fetching prefill parallel info from bootstrap: %s", exc) + return None + + +def _get_bootstrap_info_from_server(bootstrap_addr, engine_rank, target_dp_group): + """Fetch the bootstrap info from the bootstrap server.""" + try: + url = f"http://{bootstrap_addr}/route?engine_rank={engine_rank}&target_dp_group={target_dp_group}" + response = requests.get(url, timeout=5) + if response.status_code == 200: + bootstrap_info = response.json() + return bootstrap_info + else: + logger.error( + "Failed to get prefill server info: %s, %s", + response.status_code, + response.text, + ) + return None + except Exception as exc: + logger.error("Error fetching prefill info from bootstrap: %s", exc) + return None + + +@dataclass(frozen=True) +class ReceiverRoutePlan: + target_tp_rank: int | None + target_tp_ranks: tuple[int, ...] + required_prefill_response_num: int + default_required_dst_info_num: int + transfer_plan: RankTransferPlan | None = None + supports_remote_spec_candidates: bool = True + + def required_dst_info_num_for_tp_rank(self, tp_rank: int) -> int: + if self.transfer_plan is None: + return self.default_required_dst_info_num + return self.transfer_plan.required_dst_info_num_for_prefill_rank(tp_rank) + + def fragments_for_tp_rank(self, tp_rank: int): + if self.transfer_plan is None: + return () + return self.transfer_plan.fragments_by_prefill_rank.get(tp_rank, ()) + + +def _buffer_kind_for_layer_offset( + kv_args, layer_index: int, offset_index: int +) -> BufferKind: + is_draft = layer_index >= getattr(kv_args, "target_layer_num", len(kv_args.offsets)) + if is_draft: + return BufferKind.DRAFT_K if offset_index == 0 else BufferKind.DRAFT_V + return BufferKind.TARGET_K if offset_index == 0 else BufferKind.TARGET_V + + +def _unit_lens_or_default(unit_lens, item_lens): + if unit_lens: + return tuple(int(x) for x in unit_lens) + return tuple(1 for _ in item_lens) + + +def _build_buffer_layout_pair( + *, + buffer_index: int, + buffer_kind: BufferKind, + sharded_axis: str, + prefill_item_len: int, + decode_item_len: int, + prefill_unit_len: int, + decode_unit_len: int, + prefill_tp_size: int, + decode_tp_size: int, +): + """Build compatible logical layouts for one prefill/decode buffer pair. + + Besides normal TP sharding and fully replicated buffers, this handles GQA + KV caches where prefill TP is larger than the number of distinct KV heads. + In that case multiple prefill TP ranks carry the same KV head, so the + transfer plan uses one representative rank from each replica group. + """ + + if prefill_unit_len != decode_unit_len: + raise ValueError( + f"prefill/decode unit sizes differ for {buffer_kind.value}: " + f"prefill={prefill_unit_len}, decode={decode_unit_len}" + ) + if prefill_item_len % prefill_unit_len != 0: + raise ValueError( + f"prefill item length is not unit-aligned for {buffer_kind.value}: " + f"item={prefill_item_len}, unit={prefill_unit_len}" + ) + if decode_item_len % decode_unit_len != 0: + raise ValueError( + f"decode item length is not unit-aligned for {buffer_kind.value}: " + f"item={decode_item_len}, unit={decode_unit_len}" + ) + + prefill_local_units = prefill_item_len // prefill_unit_len + decode_local_units = decode_item_len // decode_unit_len + prefill_global_units = prefill_local_units * prefill_tp_size + decode_global_units = decode_local_units * decode_tp_size + prefill_tp_replica_group_size = 1 + if prefill_global_units == decode_global_units: + logical_axis = sharded_axis + logical_size = decode_global_units + elif prefill_item_len == decode_item_len: + logical_axis = "replicated" + logical_size = decode_local_units + elif ( + sharded_axis == "kv_head" + and decode_global_units % prefill_local_units == 0 + and decode_global_units // prefill_local_units <= prefill_tp_size + and prefill_tp_size % (decode_global_units // prefill_local_units) == 0 + ): + logical_axis = sharded_axis + logical_size = decode_global_units + prefill_distinct_tp_size = decode_global_units // prefill_local_units + prefill_tp_replica_group_size = prefill_tp_size // prefill_distinct_tp_size + else: + raise ValueError( + f"unsupported heterogeneous TP buffer layout for {buffer_kind.value}: " + f"prefill_item={prefill_item_len}, decode_item={decode_item_len}, " + f"prefill_tp={prefill_tp_size}, decode_tp={decode_tp_size}, " + f"unit={decode_unit_len}" + ) + + return ( + BufferLayout( + buffer_index=buffer_index, + buffer_kind=buffer_kind, + logical_axis=logical_axis, + logical_size=logical_size, + page_size=1, + bytes_per_logical_unit=decode_unit_len, + item_stride_bytes=prefill_item_len, + tp_replica_group_size=prefill_tp_replica_group_size, + ), + BufferLayout( + buffer_index=buffer_index, + buffer_kind=buffer_kind, + logical_axis=logical_axis, + logical_size=logical_size, + page_size=1, + bytes_per_logical_unit=decode_unit_len, + item_stride_bytes=decode_item_len, + ), + ) + + +def _build_kv_buffer_layouts( + kv_args, prefill_parallel_info: PrefillParallelInfo, decode_tp_size: int +): + prefill_tp_size = prefill_parallel_info.prefill_tp_size_per_dp_rank + prefill_kv_item_lens = tuple(prefill_parallel_info.kv_item_lens) + if not prefill_kv_item_lens: + prefill_kv_item_lens = tuple( + int(x) * decode_tp_size // prefill_tp_size for x in kv_args.kv_item_lens + ) + prefill_kv_unit_lens = _unit_lens_or_default( + prefill_parallel_info.kv_unit_lens, prefill_kv_item_lens + ) + decode_kv_unit_lens = _unit_lens_or_default( + getattr(kv_args, "kv_unit_lens", []), kv_args.kv_item_lens + ) + + prefill_buffers = [] + decode_buffers = [] + for layer_index, ptr_offsets in enumerate(kv_args.offsets): + for offset_index, ptr_offset in enumerate(ptr_offsets): + buffer_kind = _buffer_kind_for_layer_offset( + kv_args, layer_index, offset_index + ) + prefill_buffer, decode_buffer = _build_buffer_layout_pair( + buffer_index=ptr_offset, + buffer_kind=buffer_kind, + sharded_axis="kv_head", + prefill_item_len=int(prefill_kv_item_lens[ptr_offset]), + decode_item_len=int(kv_args.kv_item_lens[ptr_offset]), + prefill_unit_len=int(prefill_kv_unit_lens[ptr_offset]), + decode_unit_len=int(decode_kv_unit_lens[ptr_offset]), + prefill_tp_size=prefill_tp_size, + decode_tp_size=decode_tp_size, + ) + prefill_buffers.append(prefill_buffer) + decode_buffers.append(decode_buffer) + + decode_state_item_lens = tuple(getattr(kv_args, "state_item_lens", []) or []) + prefill_state_item_lens = tuple(prefill_parallel_info.state_item_lens) + if not prefill_state_item_lens: + prefill_state_item_lens = tuple( + int(x) * decode_tp_size // prefill_tp_size for x in decode_state_item_lens + ) + prefill_state_unit_lens = _unit_lens_or_default( + prefill_parallel_info.state_unit_lens, prefill_state_item_lens + ) + decode_state_unit_lens = _unit_lens_or_default( + getattr(kv_args, "state_unit_lens", []), decode_state_item_lens + ) + + for state_index, decode_item_len in enumerate(decode_state_item_lens): + prefill_buffer, decode_buffer = _build_buffer_layout_pair( + buffer_index=state_index, + buffer_kind=BufferKind.MAMBA_STATE, + sharded_axis="state_channel", + prefill_item_len=int(prefill_state_item_lens[state_index]), + decode_item_len=int(decode_item_len), + prefill_unit_len=int(prefill_state_unit_lens[state_index]), + decode_unit_len=int(decode_state_unit_lens[state_index]), + prefill_tp_size=prefill_tp_size, + decode_tp_size=decode_tp_size, + ) + prefill_buffers.append(prefill_buffer) + decode_buffers.append(decode_buffer) + return tuple(prefill_buffers), tuple(decode_buffers) + + +def _build_non_mla_route_plan(kv_mgr, prefill_parallel_info: PrefillParallelInfo): + prefill_tp_size = prefill_parallel_info.prefill_tp_size_per_dp_rank + decode_tp_size = kv_mgr.world_size // kv_mgr.dp_size + decode_tp_rank = kv_mgr.kv_args.engine_rank % decode_tp_size + prefill_buffers, decode_buffers = _build_kv_buffer_layouts( + kv_mgr.kv_args, + prefill_parallel_info, + decode_tp_size, + ) + planner = PDTransferPlanner( + prefill_layout=ParallelLayout( + role="prefill", + world_size=prefill_tp_size, + dp_size=1, + ), + decode_layout=ParallelLayout( + role="decode", + world_size=decode_tp_size, + dp_size=1, + ), + prefill_buffers=prefill_buffers, + decode_buffers=decode_buffers, + ) + transfer_plan = planner.plan_for_decode_rank(decode_tp_rank) + target_tp_ranks = tuple(transfer_plan.target_prefill_ranks) + target_tp_rank = ( + target_tp_ranks[0] if transfer_plan.plan_kind == "identity" else None + ) + default_required_dst_info_num = ( + transfer_plan.required_dst_info_num_for_prefill_rank(target_tp_ranks[0]) + if target_tp_ranks + else 1 + ) + return ReceiverRoutePlan( + target_tp_rank=target_tp_rank, + target_tp_ranks=target_tp_ranks, + required_prefill_response_num=transfer_plan.required_prefill_response_num, + default_required_dst_info_num=default_required_dst_info_num, + transfer_plan=transfer_plan, + supports_remote_spec_candidates=transfer_plan.plan_kind == "identity", + ) + + +def _legacy_mla_route_plan( + *, + target_tp_rank: int | None, + target_tp_ranks, + required_dst_info_num: int, + required_prefill_response_num: int, +) -> ReceiverRoutePlan: + return ReceiverRoutePlan( + target_tp_rank=target_tp_rank, + target_tp_ranks=tuple(target_tp_ranks), + required_prefill_response_num=required_prefill_response_num, + default_required_dst_info_num=required_dst_info_num, + ) + + +def _calc(kv_mgr, prefill_parallel_info: PrefillParallelInfo) -> ReceiverRoutePlan: + prefill_tp_size_per_dp_rank = prefill_parallel_info.prefill_tp_size_per_dp_rank + local_tp_size_per_dp_rank = kv_mgr.world_size // kv_mgr.dp_size + + if prefill_parallel_info.enable_mla_l1_5_cache: + if not kv_mgr.is_mla_backend: + raise RuntimeError( + "PD with MLA L1.5 cache is not yet supported for non-MLA models" + ) + return _legacy_mla_route_plan( + target_tp_rank=None, + target_tp_ranks=range(prefill_tp_size_per_dp_rank), + required_dst_info_num=local_tp_size_per_dp_rank, + required_prefill_response_num=prefill_tp_size_per_dp_rank, + ) + + if not kv_mgr.is_mla_backend: + return _build_non_mla_route_plan(kv_mgr, prefill_parallel_info) + + if local_tp_size_per_dp_rank == prefill_tp_size_per_dp_rank: + target_tp_rank = kv_mgr.kv_args.engine_rank % local_tp_size_per_dp_rank + return _legacy_mla_route_plan( + target_tp_rank=target_tp_rank, + target_tp_ranks=(target_tp_rank,), + required_dst_info_num=1, + required_prefill_response_num=1, + ) + + if local_tp_size_per_dp_rank > prefill_tp_size_per_dp_rank: + target_tp_rank = (kv_mgr.kv_args.engine_rank % local_tp_size_per_dp_rank) // ( + local_tp_size_per_dp_rank // prefill_tp_size_per_dp_rank + ) + return _legacy_mla_route_plan( + target_tp_rank=target_tp_rank, + target_tp_ranks=(target_tp_rank,), + required_dst_info_num=local_tp_size_per_dp_rank + // prefill_tp_size_per_dp_rank, + required_prefill_response_num=1, + ) + + target_tp_ranks = tuple( + range( + (kv_mgr.kv_args.engine_rank % local_tp_size_per_dp_rank) + * (prefill_tp_size_per_dp_rank // local_tp_size_per_dp_rank), + (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), + ) + ) + return _legacy_mla_route_plan( + target_tp_rank=target_tp_ranks[0], + target_tp_ranks=target_tp_ranks, + required_dst_info_num=1, + required_prefill_response_num=1, + ) + + +class MooncakeKVReceiver: + _ctx = zmq.Context() + _socket_cache = {} + _socket_locks = {} + _global_lock = threading.Lock() + + def __init__( + self, mgr: MooncakeKVManagerDecode, bootstrap_addr: str, bootstrap_room: int + ): + self.kv_mgr = mgr + self.bootstrap_addr = bootstrap_addr + self.bootstrap_room = bootstrap_room + + self.session_id = self.kv_mgr.get_session_id() + self.conclude_state = None + self.init_time = None + self.prefill_enable_mla_l1_5_cache = None + self.dst_enable_mla_l1_5_cache = False + + self.kv_mgr.update_status(self.bootstrap_room, TransferPoll.Bootstrapping) + logger.info( + "[MooncakeKVReceiver.__init__] bootstrap_addr=%s bootstrap_room=%s session_id=%s", + bootstrap_addr, + bootstrap_room, + self.session_id, + ) + + prefill_parallel_info = self._get_prefill_parallel_info() + if prefill_parallel_info 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, TransferPoll.Failed) + + route_plan = _calc(self.kv_mgr, prefill_parallel_info) + self.route_plan = route_plan + self.supports_remote_spec_candidates = ( + route_plan.supports_remote_spec_candidates + ) + self.required_dst_info_num = route_plan.default_required_dst_info_num + self.kv_mgr.required_prefill_response_num_table[self.bootstrap_room] = ( + route_plan.required_prefill_response_num + ) + target_dp_group = self.bootstrap_room % prefill_parallel_info.dp_size + target_tp_key = ",".join(str(rank) for rank in route_plan.target_tp_ranks) + bootstrap_key = f"{self.bootstrap_addr}_{target_dp_group}_{target_tp_key}" + if bootstrap_key not in self.kv_mgr.connection_pool: + bootstrap_infos = self._get_bootstrap_infos(target_dp_group, route_plan) + if bootstrap_infos is None: + 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: {target_dp_group}", + ) + self.kv_mgr.update_status(self.bootstrap_room, TransferPoll.Failed) + else: + if not bootstrap_infos: + raise RuntimeError("Could not fetch bootstrap info.") + 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] + + self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].add(self.bootstrap_room) + self.kv_mgr.update_status(self.bootstrap_room, TransferPoll.Bootstrapped) + logger.info( + "[MooncakeKVReceiver.__init__] done, status set to Bootstrapped. " + "bootstrap_room=%s bootstrap_addr=%s session_id=%s", + self.bootstrap_room, + self.bootstrap_addr, + self.session_id, + ) + + def _get_prefill_parallel_info(self): + prefill_parallel_info = self.kv_mgr.prefill_parallel_info.get( + self.bootstrap_addr + ) + + if prefill_parallel_info is not None: + return prefill_parallel_info + else: + prefill_parallel_info = _get_prefill_parallel_info_from_server( + self.bootstrap_addr + ) + + if prefill_parallel_info is None: + return None + else: + logger.debug( + "Fetch prefill parallel info from [%s]: DP size:%s, TP size:%s", + self.bootstrap_addr, + prefill_parallel_info.dp_size, + prefill_parallel_info.tp_size, + ) + self.kv_mgr.prefill_parallel_info[self.bootstrap_addr] = ( + prefill_parallel_info + ) + return prefill_parallel_info + + def _get_bootstrap_infos(self, target_dp_group, route_plan: ReceiverRoutePlan): + bootstrap_infos = [] + for _target_tp_rank in route_plan.target_tp_ranks: + bootstrap_info = _get_bootstrap_info_from_server( + self.bootstrap_addr, + _target_tp_rank, + target_dp_group, + ) + if bootstrap_info is not None: + # only support MLA for now: select one prefill rank as real rank + bootstrap_info["is_dummy"] = not bool( + _target_tp_rank == route_plan.target_tp_rank + or route_plan.target_tp_rank is None + ) + bootstrap_info["required_dst_info_num"] = ( + route_plan.required_dst_info_num_for_tp_rank(_target_tp_rank) + ) + bootstrap_info["transfer_fragments"] = route_plan.fragments_for_tp_rank( + _target_tp_rank + ) + logger.debug( + "Fetched bootstrap info: %s for DP %s TP %s", + bootstrap_info, + target_dp_group, + _target_tp_rank, + ) + bootstrap_infos.append(bootstrap_info) + else: + return None + return bootstrap_infos + + def _register_kv_args(self): + for bootstrap_info in self.bootstrap_infos: + self.prefill_server_url = ( + f"{bootstrap_info['rank_ip']}:{bootstrap_info['rank_port']}" + ) + logger.info( + "[MooncakeKVReceiver._register_kv_args] sending kv_args to prefill=%s bootstrap_room=%s session_id=%s", + self.prefill_server_url, + self.bootstrap_room, + self.session_id, + ) + packed_kv_data_ptrs = b"".join( + struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.kv_data_ptrs + ) + packed_state_data_ptrs = b"".join( + struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.state_data_ptrs + ) + + sock, lock = self._connect("tcp://" + self.prefill_server_url) + with lock: + sock.send_multipart( + [ + "None".encode("ascii"), + get_local_ip_by_remote().encode("ascii"), + str(self.kv_mgr.rank_port).encode("ascii"), + self.session_id.encode("ascii"), + packed_kv_data_ptrs, + b"", # aux_data_ptrs removed; kept as empty frame for protocol compat + packed_state_data_ptrs, + # Include decode_prefix_len for kv_args registration + str(getattr(self, "decode_prefix_len", 0)).encode("ascii"), + ] + ) + + @classmethod + def _connect(cls, endpoint: str): + with cls._global_lock: + if endpoint not in cls._socket_cache: + sock = cls._ctx.socket(zmq.PUSH) + sock.connect(endpoint) + cls._socket_cache[endpoint] = sock + cls._socket_locks[endpoint] = threading.Lock() + return cls._socket_cache[endpoint], cls._socket_locks[endpoint] + + def prefill( + self, + kv_indices: npt.NDArray[np.int64], + aux_index: int | None = None, + decode_prefix_len: int | None = 0, + mla_l1_5_args: PageTransferMetadata | None = None, + mamba_indices: npt.NDArray[np.int64] | None = None, + ): + logger.info( + "[MooncakeKVReceiver.init] bootstrap_room=%s kv_indices_len=%d aux_index=%s decode_prefix_len=%s", + self.bootstrap_room, + len(kv_indices), + aux_index, + decode_prefix_len, + ) + # Store decode_prefix_len to be sent back to prefill + self.decode_prefix_len = decode_prefix_len + dst_page_transfer_mask = None + dst_page_local_indices = None + if mla_l1_5_args is not None: + dst_page_transfer_mask = mla_l1_5_args.page_transfer_mask + dst_page_local_indices = mla_l1_5_args.page_local_indices + + for bootstrap_info in self.bootstrap_infos: + self.prefill_server_url = ( + f"{bootstrap_info['rank_ip']}:{bootstrap_info['rank_port']}" + ) + is_dummy = bootstrap_info["is_dummy"] + + logger.info( + "[MooncakeKVReceiver.init] sending pre-alloc multipart to prefill=%s bootstrap_room=%s is_dummy=%s", + self.prefill_server_url, + self.bootstrap_room, + bootstrap_info["is_dummy"], + ) + sock, lock = self._connect("tcp://" + self.prefill_server_url) + with lock: + message_parts = [ + str(self.bootstrap_room).encode("ascii"), + get_local_ip_by_remote().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( + bootstrap_info.get( + "required_dst_info_num", self.required_dst_info_num + ) + ).encode("ascii"), + # Send decode_prefix_len as additional message part + ( + str(self.decode_prefix_len).encode("ascii") + if not is_dummy + else b"" + ), + ( + str(int(self.dst_enable_mla_l1_5_cache)).encode("ascii") + if not is_dummy + else b"" + ), + ( + dst_page_transfer_mask.tobytes() + if (not is_dummy and dst_page_transfer_mask is not None) + else b"" + ), + ( + dst_page_local_indices.tobytes() + if (not is_dummy and dst_page_local_indices is not None) + else b"" + ), + ( + mamba_indices.tobytes() + if (not is_dummy and mamba_indices is not None) + else b"" + ), + ] + transfer_fragments = bootstrap_info.get("transfer_fragments", ()) + if not is_dummy and transfer_fragments: + message_parts.extend(encode_transfer_fragments(transfer_fragments)) + sock.send_multipart(message_parts) + self.init_time = time.time() + + def poll(self) -> TransferPoll: + if self.conclude_state is None: + status = self.kv_mgr.check_status(self.bootstrap_room) + if status in (TransferPoll.Success, TransferPoll.Failed): + self.conclude_state = status + elif status == TransferPoll.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 TOKENSPEED_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 TransferPoll.WaitingForInput", + ) + self.conclude_state = TransferPoll.Failed + return TransferPoll.Failed + elif status == TransferPoll.Transferring: + logger.warning( + "Req(room=%s) in Transferring, which is unexpected", + self.bootstrap_room, + ) + + 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 = TransferPoll.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, self.bootstrap_addr) diff --git a/python/tokenspeed/runtime/pd/mooncake/sender.py b/python/tokenspeed/runtime/pd/mooncake/sender.py new file mode 100644 index 0000000..afb8e20 --- /dev/null +++ b/python/tokenspeed/runtime/pd/mooncake/sender.py @@ -0,0 +1,200 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import time + +import numpy as np +import numpy.typing as npt + +from tokenspeed.runtime.pd.base.status import TransferPoll +from tokenspeed.runtime.pd.mooncake.entities import KVTransferError +from tokenspeed.runtime.pd.utils import PageTransferMetadata +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +class MooncakeKVSender: + def __init__( + self, + mgr, # MooncakeKVManagerPrefill + bootstrap_addr: str, + bootstrap_room: int, + ): + self.kv_mgr = mgr + self.bootstrap_server_url = bootstrap_addr + self.bootstrap_room = bootstrap_room + self.kv_mgr.update_status(bootstrap_room, TransferPoll.Bootstrapping) + logger.info( + "[MooncakeKVSender.__init__] bootstrap_room=%s bootstrap_addr=%s status=Bootstrapping", + bootstrap_room, + bootstrap_addr, + ) + + # inner state + self.init_time = None + self.conclude_state = None + self.curr_idx = 0 + self._layerwise_transfer_started = False + + def has_layerwise_transfer(self) -> bool: + return self._layerwise_transfer_started + + def send( + self, + kv_indices: npt.NDArray[np.int64], + aux_index, + is_last, + mla_l1_5_args: PageTransferMetadata | None = None, + bootstrap_token: int = -1, + spec_candidate_ids: list[int] | None = None, + mamba_indices: npt.NDArray[np.int64] | None = None, + ): + """ + Send the kv cache at the given kv indices to the decoder server + mla_l1_5_args: optional (page_transfer_mask, page_local_indices) + page_transfer_mask: boolean mask to select decode pages that will receive data from this prefill rank + page_local_indices: remapped local page indices that this prefill rank will send + bootstrap_token: first output token produced by prefill (shipped via ZMQ status msg). + """ + index_slice = slice(self.curr_idx, self.curr_idx + len(kv_indices)) + self.curr_idx += len(kv_indices) + + logger.info( + "[MooncakeKVSender.send] bootstrap_room=%s kv_indices_len=%d is_last=%s curr_idx=%d bootstrap_token=%s", + self.bootstrap_room, + len(kv_indices), + is_last, + self.curr_idx, + bootstrap_token, + ) + + if not is_last: + self.kv_mgr.add_transfer_request( + self.bootstrap_room, + kv_indices, + index_slice, + False, + mla_l1_5_args=mla_l1_5_args, + mamba_indices=mamba_indices, + ) + else: + self.kv_mgr.add_transfer_request( + self.bootstrap_room, + kv_indices, + index_slice, + True, + aux_index=aux_index, + mla_l1_5_args=mla_l1_5_args, + bootstrap_token=bootstrap_token, + spec_candidate_ids=spec_candidate_ids, + mamba_indices=mamba_indices, + ) + + def send_layerwise( + self, + kv_indices: npt.NDArray[np.int64], + index_slice: slice, + aux_index, + is_last, + begin_cache_step: int, + layerwise_interval: int, + mla_l1_5_args: PageTransferMetadata | None = None, + bootstrap_token: int = -1, + wait_for_bootstrap_token: bool = False, + spec_candidate_ids: list[int] | None = None, + mamba_indices: npt.NDArray[np.int64] | None = None, + ): + self._layerwise_transfer_started = True + self.curr_idx = max(self.curr_idx, index_slice.stop) + + if len(kv_indices) == 0 and not is_last: + return + + logger.info( + "[MooncakeKVSender.send_layerwise] bootstrap_room=%s kv_indices_len=%d " + "slice=(%s,%s) is_last=%s begin_cache_step=%s interval=%s", + self.bootstrap_room, + len(kv_indices), + index_slice.start, + index_slice.stop, + is_last, + begin_cache_step, + layerwise_interval, + ) + self.kv_mgr.add_transfer_request( + self.bootstrap_room, + kv_indices, + index_slice, + is_last, + aux_index=aux_index if is_last else None, + mla_l1_5_args=mla_l1_5_args, + bootstrap_token=bootstrap_token, + begin_cache_step=begin_cache_step, + layerwise_interval=layerwise_interval, + wait_for_bootstrap_token=wait_for_bootstrap_token, + spec_candidate_ids=spec_candidate_ids, + mamba_indices=mamba_indices, + ) + + def poll(self) -> TransferPoll: + if self.conclude_state is None: + status = self.kv_mgr.check_status(self.bootstrap_room) + if status in (TransferPoll.Success, TransferPoll.Failed): + self.conclude_state = status + elif status == TransferPoll.Bootstrapping: + if self.init_time is not None: + now = time.time() + elapsed = now - self.init_time + if elapsed >= self.kv_mgr.bootstrap_time_out: + 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 TOKENSPEED_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 TransferPoll.Bootstrapping", + ) + self.conclude_state = TransferPoll.Failed + return TransferPoll.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 = TransferPoll.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, self.bootstrap_server_url + ) diff --git a/python/tokenspeed/runtime/pd/prefill_executor.py b/python/tokenspeed/runtime/pd/prefill_executor.py new file mode 100644 index 0000000..edcd2be --- /dev/null +++ b/python/tokenspeed/runtime/pd/prefill_executor.py @@ -0,0 +1,333 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import numpy as np + +from tokenspeed.runtime.pd.base.bootstrap import BootstrapInfo +from tokenspeed.runtime.pd.base.status import TransferPoll +from tokenspeed.runtime.pd.mooncake.prefill import ( + MooncakeKVManagerPrefill, + MooncakeKVSender, +) +from tokenspeed.runtime.pd.utils import ( + TransferBackend, + poll_and_all_reduce, +) +from tokenspeed.runtime.utils import get_colorful_logger +from tokenspeed.runtime.utils.dispatch import TypeBasedDispatcher + +logger = get_colorful_logger(__name__) + +from tokenspeed_scheduler import PD, Forward + + +class DisaggPrefillExecutor: + def __init__( + self, backend: TransferBackend, args, kv_args, gloo_group, page_size: int + ): + self.transfer_backend = backend + self.bootstrap_port = args.bootstrap_port + self.page_size = page_size + self._dispatcher = TypeBasedDispatcher( + [ + (Forward.FlatForwardOp, self._decode), + ] + ) + self.senders: dict[int, MooncakeKVSender] = {} + self.kv_manager = MooncakeKVManagerPrefill(args, kv_args) + self.gloo_group = gloo_group + self._local_states = {} + self._layerwise_enabled = False + self._layerwise_interval = 1 + # request_id -> bootstrap metadata, populated after the prefill forward pass. + # Request ids and bootstrap rooms are stable across request-pool slot reuse. + self._request_token: dict[str, int] = {} + self._request_spec_candidate_ids: dict[str, list[int]] = {} + self._layerwise_token_published = set() + + def store_prefill_token( + self, + request_id: str, + aux_index: int, + token: int, + spec_candidate_ids: list[int] | None = None, + ) -> None: + """Called by event_loop after prefill forward to record the first output token.""" + self._request_token[request_id] = token + if spec_candidate_ids is not None: + self._request_spec_candidate_ids[request_id] = spec_candidate_ids + if self._layerwise_enabled: + sender = self.senders.get(request_id) + if sender is None: + logger.warning( + "Prefill token arrived before sender registration for request_id=%s", + request_id, + ) + return + self.kv_manager.set_prefill_metadata( + sender.bootstrap_room, + token, + spec_candidate_ids, + ) + self._layerwise_token_published.add(request_id) + + def register_layerwise_step_counter(self, step_counter, interval: int) -> None: + self._layerwise_enabled = True + self._layerwise_interval = max(int(interval), 1) + self.kv_manager.register_layerwise_step_counter( + step_counter, self._layerwise_interval + ) + + def _bootstrap(self, request_id, info): + self.senders[request_id] = MooncakeKVSender( + mgr=self.kv_manager, + bootstrap_addr=f"{info.bootstrap_host}:{info.bootstrap_port}", + bootstrap_room=info.bootstrap_room, + ) + + def _drop_request_state(self, req_id: str) -> None: + # Best-effort cleanup of all per-request state so failed/aborted + # requests do not leak into the bookkeeping dicts. request_id is + # stable (not a reusable slot index), so without explicit pop here + # these entries would live until the engine restarts. + sender = self.senders.pop(req_id, None) + if sender is not None: + self.kv_manager.discard_expired_metadata_room(sender.bootstrap_room) + self._local_states.pop(req_id, None) + self._request_token.pop(req_id, None) + self._request_spec_candidate_ids.pop(req_id, None) + self._layerwise_token_published.discard(req_id) + + @staticmethod + def _mamba_indices(op, index: int): + indices = getattr(op, "mamba_pool_indices", None) + if indices is None or index >= len(indices): + return None + slot = int(indices[index]) + if slot < 0: + return None + return np.array([slot], dtype=np.int64) + + @staticmethod + def _mamba_optional_index(op, attr: str, index: int): + indices = getattr(op, attr, None) + if indices is None or index >= len(indices): + return None + slot = int(indices[index]) + if slot < 0: + return None + return slot + + @classmethod + def _mamba_transfer_indices(cls, op, index: int): + working = cls._mamba_indices(op, index) + if working is None: + return None + + slots = [int(x) for x in working.tolist()] + checkpoint_src = cls._mamba_optional_index( + op, "mamba_checkpoint_dst_indices", index + ) + if checkpoint_src is not None and checkpoint_src not in slots: + slots.append(checkpoint_src) + return np.array(slots, dtype=np.int64) + + def _decode_prefix_len(self, bootstrap_room: int) -> int: + transfer_info = next( + t + for t in self.kv_manager.transfer_infos[bootstrap_room].values() + if not t.is_dummy + ) + return transfer_info.decode_prefix_len + + def _prefill_page_window(self, op, index: int, sender): + decode_prefix_len = self._decode_prefix_len(sender.bootstrap_room) + if decode_prefix_len % self.page_size != 0: + raise ValueError( + "decode_prefix_len must be divisible by page_size: " + f"{decode_prefix_len=} {self.page_size=}" + ) + + chunk_begin = op.extend_prefix_lens[index] + chunk_end = chunk_begin + op.input_lengths[index] + is_last = chunk_end >= op.prefill_lengths[index] + + decode_prefix_pages = decode_prefix_len // self.page_size + start_page = max( + decode_prefix_pages + sender.curr_idx, + chunk_begin // self.page_size, + ) + if is_last: + end_page = (chunk_end + self.page_size - 1) // self.page_size + else: + end_page = chunk_end // self.page_size + end_page = min(end_page, len(op.occupied_pages[index])) + start_page = min(start_page, end_page) + + index_slice = slice( + start_page - decode_prefix_pages, + end_page - decode_prefix_pages, + ) + return ( + np.array(op.occupied_pages[index][start_page:end_page], dtype=np.int64), + index_slice, + is_last, + ) + + def prepare_prefill(self, op) -> None: + if not self._layerwise_enabled or op.num_extends() == 0: + return + begin_cache_step = self.kv_manager.reserve_layerwise_cache_steps() + for i, request_id in enumerate(op.request_ids[: op.num_extends()]): + sender = self.senders.get(request_id) + if sender is None: + logger.debug( + "[prefill][prepare_prefill] skipping request_id=%s without sender", + request_id, + ) + self._drop_request_state(request_id) + continue + kv_indices, index_slice, is_last = self._prefill_page_window(op, i, sender) + if len(kv_indices) == 0 and not is_last: + continue + mamba_indices = self._mamba_transfer_indices(op, i) if is_last else None + sender.send_layerwise( + kv_indices, + index_slice, + op.request_pool_indices[i], + is_last, + begin_cache_step=begin_cache_step, + layerwise_interval=self._layerwise_interval, + wait_for_bootstrap_token=is_last, + mamba_indices=mamba_indices, + ) + + def _decode(self, op): + is_last = True + + for i, request_id in enumerate(op.request_ids): + sender = self.senders.get(request_id) + if sender is None: + logger.debug( + "[prefill][_decode] skipping request_id=%s without sender", + request_id, + ) + self._drop_request_state(request_id) + continue + aux_index = op.request_pool_indices[i] + bootstrap_token = self._request_token.pop(request_id, -1) + spec_candidate_ids = self._request_spec_candidate_ids.pop(request_id, None) + if sender.has_layerwise_transfer(): + if request_id not in self._layerwise_token_published: + self.kv_manager.set_prefill_metadata( + sender.bootstrap_room, + bootstrap_token, + spec_candidate_ids, + ) + self._layerwise_token_published.discard(request_id) + continue + + bootstrap_room = sender.bootstrap_room + decode_prefix_len = self._decode_prefix_len(bootstrap_room) + if decode_prefix_len % self.page_size != 0: + raise ValueError( + "decode_prefix_len must be divisible by page_size: " + f"{decode_prefix_len=} {self.page_size=}" + ) + kv_indices = np.array( + op.occupied_pages[i][decode_prefix_len // self.page_size :], + dtype=np.int64, + ) + mamba_indices = self._mamba_transfer_indices(op, i) + logger.debug( + "[prefill][_decode] rid=%s aux_index=%d kv_indices(len=%d)=%s bootstrap_token=%d", + request_id, + aux_index, + len(kv_indices), + kv_indices, + bootstrap_token, + ) + sender.send( + kv_indices, + aux_index, + is_last, + bootstrap_token=bootstrap_token, + spec_candidate_ids=spec_candidate_ids, + mamba_indices=mamba_indices, + ) + + def register(self, request_id: str, bootstrap_info: BootstrapInfo): + self._local_states[request_id] = TransferPoll.Bootstrapping + self._bootstrap(request_id, bootstrap_info) + + def abort(self, request_id: str, bootstrap_info: BootstrapInfo) -> None: + """EPD: the prefill aborted this request before registering a KV sender + (embedding receive timed out). Signal the dual-dispatched decode so its KV + receiver fails instead of waiting forever. No sender was registered, so + there is nothing to tear down on this side.""" + self.kv_manager.abort_room( + bootstrap_info.bootstrap_room, + f"EPD: prefill aborted request {request_id} (embedding receive timed out)", + ) + + def execute(self, op): + self._dispatcher(op) + + def generate_events(self): + if not self.senders: + return [] + polls = poll_and_all_reduce(self.senders.values(), self.gloo_group) + + events = [] + to_remove = [] + for req_id, poll in zip(list(self.senders.keys()), polls): + if ( + self._local_states[req_id] == TransferPoll.Bootstrapping + and poll == TransferPoll.Bootstrapped + ): + logger.debug( + "[prefill][generate_events] rid=%s -> BootstrappedEvent", req_id + ) + events.append(PD.BootstrappedEvent(req_id)) + self._local_states[req_id] = TransferPoll.Bootstrapped + elif poll == TransferPoll.Failed: + logger.warning( + "[prefill][generate_events] rid=%s -> FailedEvent", req_id + ) + events.append(PD.FailedEvent(req_id)) + to_remove.append(req_id) + elif ( + self._local_states[req_id] == TransferPoll.Bootstrapped + and poll == TransferPoll.Success + ): + self._local_states[req_id] = TransferPoll.Success + logger.debug( + "[prefill][generate_events] rid=%s -> SucceededEvent", req_id + ) + events.append(PD.SucceededEvent(req_id)) + to_remove.append(req_id) + else: + pass + for req_id in to_remove: + self._drop_request_state(req_id) + + return events diff --git a/python/tokenspeed/runtime/pd/transfer_plan.py b/python/tokenspeed/runtime/pd/transfer_plan.py new file mode 100644 index 0000000..54478d3 --- /dev/null +++ b/python/tokenspeed/runtime/pd/transfer_plan.py @@ -0,0 +1,467 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import json +from dataclasses import dataclass +from enum import Enum +from typing import Literal + + +class UnsupportedPDLayoutError(ValueError): + pass + + +class BufferKind(str, Enum): + TARGET_K = "target_k" + TARGET_V = "target_v" + DRAFT_K = "draft_k" + DRAFT_V = "draft_v" + MAMBA_STATE = "mamba_state" + + +@dataclass(frozen=True) +class ParallelLayout: + role: Literal["prefill", "decode"] + world_size: int + dp_size: int = 1 + + def __post_init__(self): + if self.world_size <= 0: + raise UnsupportedPDLayoutError("world_size must be positive") + if self.dp_size <= 0: + raise UnsupportedPDLayoutError("dp_size must be positive") + if self.world_size % self.dp_size != 0: + raise UnsupportedPDLayoutError( + f"world_size={self.world_size} must be divisible by dp_size={self.dp_size}" + ) + + @property + def tp_size_per_dp(self) -> int: + return self.world_size // self.dp_size + + +@dataclass(frozen=True) +class BufferLayout: + """Logical layout for one cache/state buffer. + + ``tp_replica_group_size`` describes TP ranks that hold the same logical + shard. It is used by GQA/MQA-style KV caches when the prefill TP size is + larger than the number of distinct KV heads. + """ + + buffer_index: int + buffer_kind: BufferKind + logical_axis: Literal["kv_head", "state_channel", "replicated"] + logical_size: int + page_size: int + bytes_per_logical_unit: int + item_stride_bytes: int + tp_replica_group_size: int = 1 + + def __post_init__(self): + if self.logical_size <= 0: + raise UnsupportedPDLayoutError("logical_size must be positive") + if self.page_size <= 0: + raise UnsupportedPDLayoutError("page_size must be positive") + if self.bytes_per_logical_unit <= 0: + raise UnsupportedPDLayoutError("bytes_per_logical_unit must be positive") + if self.item_stride_bytes <= 0: + raise UnsupportedPDLayoutError("item_stride_bytes must be positive") + if self.tp_replica_group_size <= 0: + raise UnsupportedPDLayoutError("tp_replica_group_size must be positive") + + +@dataclass(frozen=True) +class TransferFragment: + buffer_index: int + buffer_kind: BufferKind + src_rank: int + dst_rank: int + src_page_stride_bytes: int + dst_page_stride_bytes: int + src_byte_offset: int + dst_byte_offset: int + bytes_per_page: int + page_count: int | None = None + + +TRANSFER_PLAN_PROTOCOL_VERSION = 1 + + +def encode_transfer_fragments( + fragments: tuple[TransferFragment, ...], +) -> tuple[bytes, bytes]: + payload = [ + { + "buffer_index": fragment.buffer_index, + "buffer_kind": fragment.buffer_kind.value, + "src_rank": fragment.src_rank, + "dst_rank": fragment.dst_rank, + "src_page_stride_bytes": fragment.src_page_stride_bytes, + "dst_page_stride_bytes": fragment.dst_page_stride_bytes, + "src_byte_offset": fragment.src_byte_offset, + "dst_byte_offset": fragment.dst_byte_offset, + "bytes_per_page": fragment.bytes_per_page, + "page_count": fragment.page_count, + } + for fragment in fragments + ] + return ( + str(TRANSFER_PLAN_PROTOCOL_VERSION).encode("ascii"), + json.dumps(payload, separators=(",", ":")).encode("utf-8"), + ) + + +def decode_transfer_fragments( + version_frame: bytes | None, + payload_frame: bytes | None, +) -> tuple[TransferFragment, ...]: + if not version_frame and not payload_frame: + return () + if not version_frame or not payload_frame: + raise UnsupportedPDLayoutError("incomplete transfer plan frames") + + try: + version = int(version_frame.decode("ascii")) + except ValueError as exc: + raise UnsupportedPDLayoutError( + "invalid transfer plan protocol version" + ) from exc + if version != TRANSFER_PLAN_PROTOCOL_VERSION: + raise UnsupportedPDLayoutError( + f"unsupported transfer plan protocol version={version}" + ) + + raw_fragments = json.loads(payload_frame.decode("utf-8")) + return tuple( + TransferFragment( + buffer_index=int(fragment["buffer_index"]), + buffer_kind=BufferKind(fragment["buffer_kind"]), + src_rank=int(fragment["src_rank"]), + dst_rank=int(fragment["dst_rank"]), + src_page_stride_bytes=int(fragment["src_page_stride_bytes"]), + dst_page_stride_bytes=int(fragment["dst_page_stride_bytes"]), + src_byte_offset=int(fragment["src_byte_offset"]), + dst_byte_offset=int(fragment["dst_byte_offset"]), + bytes_per_page=int(fragment["bytes_per_page"]), + page_count=( + None if fragment["page_count"] is None else int(fragment["page_count"]) + ), + ) + for fragment in raw_fragments + ) + + +@dataclass(frozen=True) +class RankTransferPlan: + plan_kind: Literal["identity", "fragmented"] + target_dp_group: int + target_prefill_ranks: tuple[int, ...] + required_prefill_response_num: int + fragments_by_prefill_rank: dict[int, tuple[TransferFragment, ...]] + required_dst_info_num_by_prefill_rank: dict[int, int] + + def required_dst_info_num_for_prefill_rank(self, prefill_rank: int) -> int: + return self.required_dst_info_num_by_prefill_rank[prefill_rank] + + +@dataclass(frozen=True) +class _Interval: + start: int + end: int + + @property + def length(self) -> int: + return self.end - self.start + + def intersect(self, other: "_Interval") -> "_Interval | None": + start = max(self.start, other.start) + end = min(self.end, other.end) + if start >= end: + return None + return _Interval(start, end) + + +class PDTransferPlanner: + def __init__( + self, + *, + prefill_layout: ParallelLayout, + decode_layout: ParallelLayout, + prefill_buffers: tuple[BufferLayout, ...], + decode_buffers: tuple[BufferLayout, ...], + ): + self.prefill_layout = prefill_layout + self.decode_layout = decode_layout + self.prefill_buffers = prefill_buffers + self.decode_buffers = decode_buffers + self._validate_buffers() + self._validate_alignment() + self._required_dst_info_num_by_prefill_rank = self._calc_source_fanout() + + def plan_for_decode_rank(self, decode_rank: int) -> RankTransferPlan: + decode_tp_size = self.decode_layout.tp_size_per_dp + if decode_rank < 0 or decode_rank >= self.decode_layout.world_size: + raise UnsupportedPDLayoutError(f"decode_rank={decode_rank} is out of range") + + target_dp_group = decode_rank // decode_tp_size + decode_tp_rank = decode_rank % decode_tp_size + + if self._can_use_identity_plan() and ( + self.prefill_layout.tp_size_per_dp == decode_tp_size + ): + prefill_rank = ( + target_dp_group * self.prefill_layout.tp_size_per_dp + decode_tp_rank + ) + return RankTransferPlan( + plan_kind="identity", + target_dp_group=target_dp_group, + target_prefill_ranks=(prefill_rank,), + required_prefill_response_num=1, + fragments_by_prefill_rank={}, + required_dst_info_num_by_prefill_rank={ + prefill_rank: self._required_dst_info_num_by_prefill_rank[ + prefill_rank + ] + }, + ) + + fragments: dict[int, list[TransferFragment]] = {} + for prefill_buffer, decode_buffer in zip( + self.prefill_buffers, self.decode_buffers + ): + if prefill_buffer.logical_axis == "replicated": + prefill_tp_rank = self._replicated_source_tp_rank( + self.prefill_layout.tp_size_per_dp, + decode_tp_size, + decode_tp_rank, + ) + prefill_rank = ( + target_dp_group * self.prefill_layout.tp_size_per_dp + + prefill_tp_rank + ) + fragment = TransferFragment( + buffer_index=prefill_buffer.buffer_index, + buffer_kind=prefill_buffer.buffer_kind, + src_rank=prefill_rank, + dst_rank=decode_rank, + src_page_stride_bytes=prefill_buffer.item_stride_bytes, + dst_page_stride_bytes=decode_buffer.item_stride_bytes, + src_byte_offset=0, + dst_byte_offset=0, + bytes_per_page=decode_buffer.item_stride_bytes, + ) + fragments.setdefault(prefill_rank, []).append(fragment) + continue + + decode_interval = self._rank_interval_for_buffer( + decode_buffer, + self.decode_layout, + decode_tp_rank, + ) + if decode_interval is None: + continue + for prefill_tp_rank in range(self.prefill_layout.tp_size_per_dp): + prefill_rank = ( + target_dp_group * self.prefill_layout.tp_size_per_dp + + prefill_tp_rank + ) + prefill_interval = self._rank_interval_for_buffer( + prefill_buffer, + self.prefill_layout, + prefill_tp_rank, + ) + if prefill_interval is None: + continue + intersection = prefill_interval.intersect(decode_interval) + if intersection is None: + continue + + fragment = TransferFragment( + buffer_index=prefill_buffer.buffer_index, + buffer_kind=prefill_buffer.buffer_kind, + src_rank=prefill_rank, + dst_rank=decode_rank, + src_page_stride_bytes=prefill_buffer.item_stride_bytes, + dst_page_stride_bytes=decode_buffer.item_stride_bytes, + src_byte_offset=(intersection.start - prefill_interval.start) + * prefill_buffer.bytes_per_logical_unit, + dst_byte_offset=(intersection.start - decode_interval.start) + * decode_buffer.bytes_per_logical_unit, + bytes_per_page=intersection.length + * prefill_buffer.bytes_per_logical_unit, + ) + fragments.setdefault(prefill_rank, []).append(fragment) + + fragments_by_rank = { + rank: tuple(rank_fragments) + for rank, rank_fragments in sorted(fragments.items()) + } + target_prefill_ranks = tuple(fragments_by_rank) + return RankTransferPlan( + plan_kind="fragmented", + target_dp_group=target_dp_group, + target_prefill_ranks=target_prefill_ranks, + required_prefill_response_num=len(target_prefill_ranks), + fragments_by_prefill_rank=fragments_by_rank, + required_dst_info_num_by_prefill_rank={ + rank: self._required_dst_info_num_by_prefill_rank[rank] + for rank in target_prefill_ranks + }, + ) + + def _validate_buffers(self) -> None: + if len(self.prefill_buffers) != len(self.decode_buffers): + raise UnsupportedPDLayoutError("prefill/decode buffer counts differ") + for prefill_buffer, decode_buffer in zip( + self.prefill_buffers, self.decode_buffers + ): + if prefill_buffer.buffer_index != decode_buffer.buffer_index: + raise UnsupportedPDLayoutError("prefill/decode buffer indexes differ") + if prefill_buffer.buffer_kind != decode_buffer.buffer_kind: + raise UnsupportedPDLayoutError("prefill/decode buffer kinds differ") + if prefill_buffer.logical_axis != decode_buffer.logical_axis: + raise UnsupportedPDLayoutError("prefill/decode logical axes differ") + if prefill_buffer.logical_size != decode_buffer.logical_size: + raise UnsupportedPDLayoutError("prefill/decode logical sizes differ") + if ( + prefill_buffer.bytes_per_logical_unit + != decode_buffer.bytes_per_logical_unit + ): + raise UnsupportedPDLayoutError( + "prefill/decode logical unit sizes differ" + ) + + def _validate_alignment(self) -> None: + for layout, buffers in ( + (self.prefill_layout, self.prefill_buffers), + (self.decode_layout, self.decode_buffers), + ): + for buffer in buffers: + if buffer.logical_axis == "replicated": + continue + if layout.tp_size_per_dp % buffer.tp_replica_group_size != 0: + raise UnsupportedPDLayoutError( + "tp replica group must divide TP size for " + f"buffer_kind={buffer.buffer_kind.value}: " + f"tp_size_per_dp={layout.tp_size_per_dp}, " + f"tp_replica_group_size={buffer.tp_replica_group_size}" + ) + effective_tp_size = ( + layout.tp_size_per_dp // buffer.tp_replica_group_size + ) + if buffer.logical_size % effective_tp_size != 0: + raise UnsupportedPDLayoutError( + "non-aligned TP heterogeneous mapping for " + f"buffer_kind={buffer.buffer_kind.value}: logical_size=" + f"{buffer.logical_size}, effective_tp_size={effective_tp_size}" + ) + item_units = buffer.item_stride_bytes // buffer.bytes_per_logical_unit + required_units = buffer.logical_size // effective_tp_size + if item_units < required_units: + raise UnsupportedPDLayoutError( + "buffer item is smaller than its logical shard for " + f"buffer_kind={buffer.buffer_kind.value}: item_units=" + f"{item_units}, required_units={required_units}" + ) + + def _calc_source_fanout(self) -> dict[int, int]: + fanout = {rank: 0 for rank in range(self.prefill_layout.world_size)} + for decode_rank in range(self.decode_layout.world_size): + decode_tp_rank = decode_rank % self.decode_layout.tp_size_per_dp + target_dp_group = decode_rank // self.decode_layout.tp_size_per_dp + intersected_prefill_ranks = set() + for prefill_buffer, decode_buffer in zip( + self.prefill_buffers, self.decode_buffers + ): + if prefill_buffer.logical_axis == "replicated": + prefill_tp_rank = self._replicated_source_tp_rank( + self.prefill_layout.tp_size_per_dp, + self.decode_layout.tp_size_per_dp, + decode_tp_rank, + ) + prefill_rank = ( + target_dp_group * self.prefill_layout.tp_size_per_dp + + prefill_tp_rank + ) + intersected_prefill_ranks.add(prefill_rank) + continue + + decode_interval = self._rank_interval_for_buffer( + decode_buffer, + self.decode_layout, + decode_tp_rank, + ) + if decode_interval is None: + continue + for prefill_tp_rank in range(self.prefill_layout.tp_size_per_dp): + prefill_interval = self._rank_interval_for_buffer( + prefill_buffer, + self.prefill_layout, + prefill_tp_rank, + ) + if prefill_interval is None: + continue + if prefill_interval.intersect(decode_interval) is None: + continue + prefill_rank = ( + target_dp_group * self.prefill_layout.tp_size_per_dp + + prefill_tp_rank + ) + intersected_prefill_ranks.add(prefill_rank) + for prefill_rank in intersected_prefill_ranks: + fanout[prefill_rank] += 1 + return fanout + + def _can_use_identity_plan(self) -> bool: + return all( + prefill_buffer.tp_replica_group_size == 1 + and decode_buffer.tp_replica_group_size == 1 + for prefill_buffer, decode_buffer in zip( + self.prefill_buffers, self.decode_buffers + ) + ) + + @staticmethod + def _rank_interval(logical_size: int, tp_size: int, tp_rank: int) -> _Interval: + local_size = logical_size // tp_size + start = tp_rank * local_size + return _Interval(start, start + local_size) + + @staticmethod + def _rank_interval_for_buffer( + buffer: BufferLayout, layout: ParallelLayout, tp_rank: int + ) -> _Interval | None: + replica_group_size = buffer.tp_replica_group_size + if tp_rank % replica_group_size != 0: + return None + effective_tp_size = layout.tp_size_per_dp // replica_group_size + effective_tp_rank = tp_rank // replica_group_size + return PDTransferPlanner._rank_interval( + buffer.logical_size, effective_tp_size, effective_tp_rank + ) + + @staticmethod + def _replicated_source_tp_rank( + prefill_tp_size: int, decode_tp_size: int, decode_tp_rank: int + ) -> int: + return (decode_tp_rank * prefill_tp_size) // decode_tp_size diff --git a/python/tokenspeed/runtime/pd/utils.py b/python/tokenspeed/runtime/pd/utils.py new file mode 100755 index 0000000..b5f96ba --- /dev/null +++ b/python/tokenspeed/runtime/pd/utils.py @@ -0,0 +1,484 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Shared helpers for PD transfer runtime components.""" + +from __future__ import annotations + +import ctypes +import dataclasses +import os +import random +import threading +import warnings +from collections import deque +from enum import Enum +from typing import TYPE_CHECKING + +import numpy as np +import numpy.typing as npt +import requests +import torch +import torch.distributed as dist + +from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput +from tokenspeed.runtime.utils import get_colorful_logger +from tokenspeed.runtime.utils.network import get_ip + +if TYPE_CHECKING: + from tokenspeed.runtime.engine.request import Req + +# env var for testing failure, convert to float explicitly +FAILURE_PROB = float(os.getenv("DISAGGREGATION_TEST_FAILURE_PROB", 0)) +logger = get_colorful_logger(__name__) + + +class DisaggregationMode(Enum): + NULL = "null" + PREFILL = "prefill" + DECODE = "decode" + ENCODE = "encode" + + +class FastQueue: + class Empty(Exception): + """Exception raised when the queue is empty.""" + + pass + + def __init__(self): + self._buf = deque() + self._cond = threading.Condition() + + def put(self, item): + with self._cond: + self._buf.append(item) + self._cond.notify() + + def get(self): + with self._cond: + while not self._buf: + self._cond.wait() + return self._buf.popleft() + + def get_nowait(self): + with self._cond: + if not self._buf: + raise FastQueue.Empty() + return self._buf.popleft() + + +def poll_and_all_reduce(pollers, gloo_group): + """Poll transfer state and all-reduce the result across the gloo group.""" + # At a certain probability, mark the poll as failed to simulate failure. + if FAILURE_PROB > 0: + from tokenspeed.runtime.pd.base.status import TransferPoll + + polls = [ + ( + int(TransferPoll.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() + + +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(range(size)) + + def available_size(self) -> int: + return len(self.free_slots) + + def alloc(self) -> int | None: + if not self.free_slots: + return None + + return self.free_slots.popleft() + + def free(self, free_index: int) -> None: + self.free_slots.append(free_index) + + +class TransferBackend(Enum): + MOONCAKE = "mooncake" + MOONCAKE_ASYNC = "mooncake_async" + + +class KVClassType(Enum): + MANAGER_PREFILL = "manager_prefill" + MANAGER_DECODE = "manager_decode" + # The async backend uses one role-agnostic manager (see get_kv_class's + # MOONCAKE_ASYNC branch); the sync backend splits into the prefill/decode + # managers above. + MANAGER = "manager" + SENDER = "sender" + RECEIVER = "receiver" + BOOTSTRAP_SERVER = "bootstrap_server" + + +def get_kv_class(transfer_backend: TransferBackend, class_type: KVClassType): + if transfer_backend == TransferBackend.MOONCAKE: + from tokenspeed.runtime.pd.mooncake.conn import ( + MooncakeKVBootstrapServer, + ) + from tokenspeed.runtime.pd.mooncake.decode import ( + MooncakeKVManagerDecode, + ) + from tokenspeed.runtime.pd.mooncake.prefill import ( + MooncakeKVManagerPrefill, + ) + from tokenspeed.runtime.pd.mooncake.receiver import ( + MooncakeKVReceiver, + ) + from tokenspeed.runtime.pd.mooncake.sender import ( + MooncakeKVSender, + ) + + class_mapping = { + KVClassType.MANAGER_PREFILL: MooncakeKVManagerPrefill, + KVClassType.MANAGER_DECODE: MooncakeKVManagerDecode, + KVClassType.SENDER: MooncakeKVSender, + KVClassType.RECEIVER: (MooncakeKVReceiver), + KVClassType.BOOTSTRAP_SERVER: MooncakeKVBootstrapServer, + } + return class_mapping.get(class_type) + + if transfer_backend == TransferBackend.MOONCAKE_ASYNC: + from tokenspeed.runtime.pd.mooncake.async_conn import ( + MooncakeAsyncKVManager, + ) + from tokenspeed.runtime.pd.mooncake.conn import ( + MooncakeKVBootstrapServer, + ) + from tokenspeed.runtime.pd.mooncake.receiver import ( + MooncakeKVReceiver, + ) + from tokenspeed.runtime.pd.mooncake.sender import ( + MooncakeKVSender, + ) + + class_mapping = { + KVClassType.MANAGER: MooncakeAsyncKVManager, + KVClassType.SENDER: MooncakeKVSender, + KVClassType.RECEIVER: (MooncakeKVReceiver), + KVClassType.BOOTSTRAP_SERVER: MooncakeKVBootstrapServer, + } + return class_mapping.get(class_type) + + raise ValueError(f"Unsupported transfer backend: {transfer_backend}") + + +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 + + +@dataclasses.dataclass +class PDRegistryRequest: + """A request to register a machine itself to the LB.""" + + mode: str + registry_url: str + bootstrap_port: int | None = None + + def __post_init__(self): + if self.mode == "prefill" and self.bootstrap_port is None: + raise ValueError("Bootstrap port must be set in PREFILL mode.") + if self.mode == "decode" and self.bootstrap_port is not None: + raise ValueError("Bootstrap port must not be set in DECODE mode.") + if self.mode not in {"prefill", "decode"}: + raise ValueError( + f"Invalid mode: {self.mode}. Must be 'prefill' or 'decode'." + ) + + +def register_disaggregation_server( + mode: str, server_port: int, bootstrap_port: int, pdlb_url: str +): + pdlb_url = pdlb_url.rstrip("/") + registered_bootstrap_port = bootstrap_port if mode == "prefill" else None + registry_request = PDRegistryRequest( + mode=mode, + registry_url=f"http://{get_ip()}:{server_port}", + bootstrap_port=registered_bootstrap_port, + ) + res = requests.post( + f"{pdlb_url}/register", + json=dataclasses.asdict(registry_request), + ) + if res.status_code != 200: + warnings.warn( + f"Failed to register disaggregation server: {res.status_code} {res.text}" + ) + else: + logger.info( + "Registered disaggregation server with %s: status_code=%s text=%s", + pdlb_url, + res.status_code, + res.text, + ) + + +def is_mla_backend(target_kv_pool) -> bool: + from tokenspeed.runtime.layers.attention.kv_cache.mla import MLATokenToKVPool + + return isinstance(target_kv_pool, MLATokenToKVPool) + + +def prepare_abort(req: Req, error_message: str, status_code=None, err_type=None): + from tokenspeed.runtime.engine.request_types import ABORT_CODE, FINISH_ABORT + + if err_type is None: + err_type = ABORT_CODE.UnknownError + + # populate finish metadata and stream output + req.finished_reason = FINISH_ABORT( + error_message, status_code=status_code, err_type=err_type + ) + + 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 = [] + + +class MetadataBuffers: + def __init__(self, size: int, max_top_logprobs_num: int = 128, device: str = "cpu"): + + # We transfer the metadata of first output token to decode + # The minimal size for RDMA is 64Bytes, so we pad it to > 64Bytes + self.device = device + 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.cached_tokens = torch.zeros((size, 1), dtype=torch.int32, 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.cached_tokens.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.cached_tokens.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.cached_tokens[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.cached_tokens[idx], + ) + + def set_buf_by_batch( + self, + output_ids: torch.Tensor, + output_buffer_indices: list[int], + logits_output: LogitsProcessorOutput, + output_token_logprobs_indices: list[tuple[int, int]] | None = None, + output_top_logprobs_indices: list[tuple[int, int]] | None = None, + cached_tokens: torch.Tensor | None = None, + ): + self.output_ids[ + torch.tensor(output_buffer_indices).to(self.device, non_blocking=True), 0 + ] = output_ids + if cached_tokens is not None: + self.cached_tokens[ + torch.tensor(output_buffer_indices).to(self.device, non_blocking=True), + 0, + ] = cached_tokens + + if output_token_logprobs_indices: + for src_idx, dst_idx in output_token_logprobs_indices: + self.output_token_logprobs_val[dst_idx][0] = ( + logits_output.next_token_top_logprobs_val[src_idx] + ) + self.output_token_logprobs_idx[dst_idx][0] = ( + logits_output.next_token_top_logprobs_idx[src_idx] + ) + + if output_top_logprobs_indices: + for src_idx, dst_idx in output_top_logprobs_indices: + self.output_top_logprobs_val[dst_idx][ + : len(logits_output.next_token_top_logprobs_val[src_idx][0]) + ] = torch.tensor( + logits_output.next_token_top_logprobs_val[src_idx][0], + dtype=torch.float32, + device="cpu", + ) + self.output_top_logprobs_idx[dst_idx][ + : len(logits_output.next_token_top_logprobs_idx[src_idx][0]) + ] = torch.tensor( + logits_output.next_token_top_logprobs_idx[src_idx][0], + dtype=torch.int32, + device="cpu", + ) + + 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" + ) + + +def group_concurrent_contiguous( + src_indices: npt.NDArray[np.int64], dst_indices: npt.NDArray[np.int64] +) -> tuple[list[npt.NDArray[np.int64]], list[npt.NDArray[np.int64]]]: + """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 + + +class StepCounter: + COUNT_NUM_MAX: int = 2**62 + + @classmethod + def is_step_ready(cls, current_step: int, target_step: int) -> bool: + # because COUNT_NUM_MAX is very large, we can make sure that if diff is > COUNT_NUM_MAX / 2 means the flush is finished + # and if the current_sent_count == task_stop_count also means the flush is not finished + # so if current_sent_count != task_stop_count and diff < COUNT_NUM_MAX / 2, the flush is not finished + return ( + target_step != current_step + and (target_step + cls.COUNT_NUM_MAX - current_step) % cls.COUNT_NUM_MAX + > cls.COUNT_NUM_MAX / 2 + ) + + def __init__(self, device: str, gpu_id: int): + # utilities for cache step + self.d_ready_cache_step = torch.tensor(0, dtype=torch.int64).cuda(gpu_id) + self.h_ready_cache_step = torch.tensor(0, dtype=torch.int64, pin_memory=True) + self.cache_step: int = 0 + + # utilities for aux step + self.d_ready_aux_step = torch.tensor(0, dtype=torch.int64).cuda(gpu_id) + self.h_ready_aux_step = torch.tensor(0, dtype=torch.int64, pin_memory=True) + self.aux_step: int = 0 + + def current_step(self) -> tuple[int, int]: + return self.cache_step, self.aux_step + + def advance_step(self, delta_cache_step: int, delta_aux_step: int): + self.cache_step = (self.cache_step + delta_cache_step) % self.COUNT_NUM_MAX + self.aux_step = (self.aux_step + delta_aux_step) % self.COUNT_NUM_MAX + + def record_cache(self): + self.d_ready_cache_step = (self.d_ready_cache_step + 1) % self.COUNT_NUM_MAX + self.h_ready_cache_step.copy_(self.d_ready_cache_step, non_blocking=True) + + def record_aux(self): + self.d_ready_aux_step = (self.d_ready_aux_step + 1) % self.COUNT_NUM_MAX + self.h_ready_aux_step.copy_(self.d_ready_aux_step, non_blocking=True) + + def query_ready_cache_step(self) -> int: + return ctypes.c_int64.from_address(self.h_ready_cache_step.data_ptr()).value + + def query_ready_aux_step(self) -> int: + return ctypes.c_int64.from_address(self.h_ready_aux_step.data_ptr()).value + + +@dataclasses.dataclass +class PageTransferMetadata: + indices_are_local: bool + page_transfer_mask: npt.NDArray[np.bool_] + page_local_indices: npt.NDArray[np.int64] | None = None diff --git a/python/tokenspeed/runtime/sampling/backends/__init__.py b/python/tokenspeed/runtime/sampling/backends/__init__.py new file mode 100644 index 0000000..69e8766 --- /dev/null +++ b/python/tokenspeed/runtime/sampling/backends/__init__.py @@ -0,0 +1,22 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +# Concrete backend modules are imported lazily in sampling.registry to avoid +# circular imports (default.py imports register_backend from registry). diff --git a/python/tokenspeed/runtime/sampling/backends/base.py b/python/tokenspeed/runtime/sampling/backends/base.py new file mode 100644 index 0000000..f970325 --- /dev/null +++ b/python/tokenspeed/runtime/sampling/backends/base.py @@ -0,0 +1,285 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch +import torch.distributed as dist + +if TYPE_CHECKING: + from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput + from tokenspeed.runtime.sampling.dp_sampling_config import DpSamplingRuntimeConfig + from tokenspeed.runtime.sampling.sampling_batch_info import SamplingBatchInfo + from tokenspeed.runtime.sampling.sampling_params import SamplingParams + from tokenspeed.runtime.utils.server_args import ServerArgs + + +DEFAULT_RANDOM_SEED = 48 +CUDA_GRAPH_VARIANT_DEFAULT = "default" +SPECULATIVE_ACCEPT_THRESHOLD_SINGLE = 1.0 +SPECULATIVE_ACCEPT_THRESHOLD_ACC = 1.0 + + +@dataclass +class SamplingBackendConfig: + + enable_nan_detection: bool = False + + # Optional logprob features — OFF by default. These are checked at server + # start / graph capture time so the fast path has zero extra compute. + # Enabling any of these enlarges the captured graph footprint. + enable_output_logprobs: bool = False + + # Sizing for pre-allocated per-backend buffers (e.g. coin buffers for + # rejection sampling). Required to keep RNG out of the CUDA graph. + max_bs: int = 1 + max_draft_tokens_per_req: int = 1 + + # Sizing for backend-owned per-request state (e.g. token-count buffers + # for penalties in FlashInferFullSamplingBackend). Indexed by req_pool_idx, not + # batch row, so the data survives batch membership changes. + max_req_pool_size: int = 0 + vocab_size: int = 0 + + device: torch.device | None = None + random_seed: int = DEFAULT_RANDOM_SEED + + # Attention TP group for sampler-output broadcast (rank 0 wins). + tp_group: tuple[int, ...] | None = None + enable_tp_sync: bool = True + + @classmethod + def from_server_args( + cls, + server_args: ServerArgs, + *, + max_bs: int, + max_draft_tokens_per_req: int, + device: str, + random_seed: int = DEFAULT_RANDOM_SEED, + max_req_pool_size: int = 0, + vocab_size: int = 0, + tp_group: tuple[int, ...] | None = None, + ) -> SamplingBackendConfig: + + return cls( + enable_nan_detection=server_args.enable_nan_detection, + enable_output_logprobs=server_args.enable_output_logprobs, + max_bs=max_bs, + max_draft_tokens_per_req=max(max_draft_tokens_per_req, 1), + max_req_pool_size=max_req_pool_size, + vocab_size=vocab_size, + device=device, + random_seed=random_seed, + tp_group=tp_group, + enable_tp_sync=not server_args.disable_sampling_tp_sync, + ) + + +class SamplingBackend(ABC): + """Shared contract for single-step sampling and multi-step spec-decode verification. + + Both methods return (output_tokens, accept_lengths). For sample(), + accept_lengths is all-ones so the downstream contract matches verify(). + + Backends that need random state override prepare() to refill per-request + buffers outside of any CUDA graph capture. + + Requests asking for params a backend doesn't implement are NOT rejected; + the backend silently applies only what it supports, so all requests go + through the same captured graph. + """ + + # Subclasses that hold per-pool-idx state (scalars like temperature / + # top_k, plus large rows like _counts / _logit_bias) flip this to True + # so prepare_step() performs flip detection + _reset_slot. Stateless + # backends (greedy) leave it False and the whole prepare_step call is + # a no-op. + _HAS_POOL_STATE: bool = False + _SUPPORTS_DP_VERIFY: bool = False + + def __init__(self, config: SamplingBackendConfig) -> None: + + self.config = config + + # Sentinel of "which rid currently owns each slot from this backend's + # point of view". rid is just a comparison value here, not a lookup + # key, so this is pool-keyed state (size O(pool_rows) strings), not + # rid-keyed state. A mismatch against the incoming rid is a flip. + if self._HAS_POOL_STATE: + pool_rows = config.max_req_pool_size + 1 + self._last_rid_per_slot: list[str | None] = [None] * pool_rows + + # Resolved once; None means maybe_broadcast is a no-op. + self._tp_pg = None + self._tp_src_global_rank: int | None = None + if ( + config.enable_tp_sync + and config.tp_group is not None + and len(config.tp_group) > 1 + ): + from tokenspeed.runtime.distributed.process_group_manager import ( + process_group_manager as pg_manager, + ) + + self._tp_pg = pg_manager.get_process_group("nccl", config.tp_group) + self._tp_src_global_rank = config.tp_group[0] + + def configure_dp_sampling(self, runtime: DpSamplingRuntimeConfig) -> None: + """Configure optional DP sampling state. + + Stateless or unsupported backends ignore this; DP-capable backends + override it to initialize backend-local communication buffers. + """ + + def maybe_broadcast(self, *tensors: torch.Tensor) -> None: + """Broadcast each tensor from tp_group[0] so all attention-TP ranks + agree. No-op when sync is off or tp_size <= 1. Graph-safe.""" + if self._tp_pg is None: + return + for t in tensors: + dist.broadcast(t, src=self._tp_src_global_rank, group=self._tp_pg) + + def prepare_step( + self, + request_ids: list[str], + request_pool_indices: list[int], + sampling_params_list: list[SamplingParams], + num_tokens_per_req: int = 1, + ) -> None: + """Called once per step, outside the CUDA graph. Two jobs: + + 1. Flip detection: a slot's owning rid changed since last step + (first-use and rid-recycling look the same). Delegates to + _reset_slot which scatters all per-slot persistent state + (scalars, counts, bias, generators). + 2. Per-step dynamic refill: coin buffers, etc. Delegated to the + subclass via _prepare_step_hook. + + Stateless backends (greedy) short-circuit both phases. + """ + + if not self._HAS_POOL_STATE: + return + + assert ( + len(request_ids) == len(request_pool_indices) == len(sampling_params_list) + ), ( + f"prepare_step expects aligned per-request lists; got " + f"rids={len(request_ids)}, pool_indices={len(request_pool_indices)}, " + f"sp_list={len(sampling_params_list)}" + ) + + pool_rows = len(self._last_rid_per_slot) + for rid, pool_idx, sp in zip( + request_ids, request_pool_indices, sampling_params_list + ): + assert ( + 0 <= pool_idx < pool_rows + ), f"pool_idx {pool_idx} out of range [0, {pool_rows}) for rid={rid}" + if self._last_rid_per_slot[pool_idx] != rid: + self._reset_slot(pool_idx, sp) + self._last_rid_per_slot[pool_idx] = rid + + self._prepare_step_hook( + num_tokens_per_req=num_tokens_per_req, + bs=len(request_pool_indices), + request_pool_indices=request_pool_indices, + ) + + def prepare_capture(self, bs: int, num_tokens_per_req: int = 1) -> None: + """Per-step refill for the capture/warm-up path. No flip detection; + the backend uses its stub generator for any RNG-fed buffers so the + captured graph sees a fully-written state. + Default: no-op. + """ + self._prepare_step_hook( + num_tokens_per_req=num_tokens_per_req, + bs=bs, + request_pool_indices=None, + ) + + def cuda_graph_capture_variants(self, num_tokens_per_req: int) -> tuple[str, ...]: + """Return sampler-specific CUDA graph variants to capture.""" + return (CUDA_GRAPH_VARIANT_DEFAULT,) + + def prepare_capture_variant( + self, + bs: int, + num_tokens_per_req: int, + variant: str, + ) -> None: + if variant != CUDA_GRAPH_VARIANT_DEFAULT: + raise ValueError(f"Unsupported CUDA graph variant: {variant}") + self.prepare_capture(bs=bs, num_tokens_per_req=num_tokens_per_req) + + def cuda_graph_replay_variant(self, num_tokens_per_req: int) -> str: + return CUDA_GRAPH_VARIANT_DEFAULT + + def _prepare_step_hook( + self, + num_tokens_per_req: int, + bs: int, + request_pool_indices: list[int] | None, + ) -> None: + """Subclass hook for per-step dynamic state (coin buffers, etc). + request_pool_indices=None is the capture path; otherwise the CPU + list from forward_op.request_pool_indices. + Default: no-op.""" + + def _reset_slot(self, pool_idx: int, sp: SamplingParams) -> None: + """Scatter all per-slot persistent state for a newly-assigned slot. + Called from prepare_step on flip. Stateful backends override.""" + raise NotImplementedError + + def reset_capture_state(self) -> None: + """Clear any per-pool state that warm-up iterations may have dirtied + before CUDA graph capture. Warm-up runs sample()/verify() against + pool row 0 (see CudaGraphWrapper capture path); stateful backends + override this to zero whatever row 0 accumulates. Default: no-op.""" + + def get_packed_output_d2h( + self, + output_tokens: torch.Tensor, + output_lengths: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor] | None: + """If the backend wrote both outputs into a single contiguous GPU + buffer, return CPU views obtained from one D2H copy. Otherwise + return None and let the caller fall back to two separate D2Hs.""" + return None + + @abstractmethod + def sample( + self, + logits_output: LogitsProcessorOutput, + sampling_info: SamplingBatchInfo, + ) -> tuple[torch.Tensor, torch.Tensor]: ... + + @abstractmethod + def verify( + self, + logits_output: LogitsProcessorOutput, + sampling_info: SamplingBatchInfo, + candidates: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: ... diff --git a/python/tokenspeed/runtime/sampling/backends/flashinfer.py b/python/tokenspeed/runtime/sampling/backends/flashinfer.py new file mode 100644 index 0000000..47ef524 --- /dev/null +++ b/python/tokenspeed/runtime/sampling/backends/flashinfer.py @@ -0,0 +1,621 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from tokenspeed_kernel.ops.sampling import argmax as sampling_argmax +from tokenspeed_kernel.ops.sampling.cuda import ( + chain_speculative_sampling_target_only, + fused_topk_topp_prepare, + fused_topk_topp_renorm, + verify_chain_greedy, +) +from tokenspeed_kernel.ops.sampling.flashinfer import ( + softmax, + top_k_renorm_prob, + top_k_top_p_sampling_from_probs, + top_p_renorm_prob, +) +from tokenspeed_kernel.ops.sampling.triton import gather_and_expand_scalars +from tokenspeed_kernel.platform import current_platform +from tokenspeed_kernel.torch_compile import get_compiler_backend + +# Resolved once at import: the fused top-k + top-p kernel is NVIDIA-only. +# On non-NVIDIA platforms (e.g. ROCm) we fall back to the back-to-back +# flashinfer renorm calls. Defining this at module scope keeps the hot path +# branch-free in the captured graph. +_FUSED_TOPK_TOPP_AVAILABLE = current_platform().is_nvidia + +from tokenspeed.runtime.distributed.dp_sampling_comm import DpSamplingComm +from tokenspeed.runtime.sampling.backends.base import ( + SPECULATIVE_ACCEPT_THRESHOLD_ACC, + SPECULATIVE_ACCEPT_THRESHOLD_SINGLE, + SamplingBackend, + SamplingBackendConfig, +) +from tokenspeed.runtime.sampling.dp_sampling_config import ( + DpSamplingRuntimeConfig, + slice_dp_vocab_mask, +) +from tokenspeed.runtime.sampling.registry import register_backend +from tokenspeed.runtime.sampling.utils import ( + coin_eps, + gather_token_logprobs_torch, +) +from tokenspeed.runtime.utils.nvtx import nvtx_range +from tokenspeed.runtime.utils.pdl import pdl_enabled + +if TYPE_CHECKING: + from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput + from tokenspeed.runtime.sampling.sampling_batch_info import SamplingBatchInfo + from tokenspeed.runtime.sampling.sampling_params import SamplingParams + + +class FlashInferSamplingBackend(SamplingBackend): + """Fast backend: fused softmax(temperature) + top_k_top_p_sampling_from_probs + for stochastic single-step sampling; cuda chain kernels (greedy + + rejection) for multi-step verification. + + Scope is deliberately narrow — temperature / top_k / top_p only — + keeping the hot path to 2 kernels. Requests asking for min_p, penalties, + or logit_bias are silently ignored; use `flashinfer_full` if any of those + matter for the workload. + """ + + _HAS_POOL_STATE = True + _SUPPORTS_DP_VERIFY = True + + def __init__(self, config: SamplingBackendConfig) -> None: + + super().__init__(config) + self._init_dp_sampling(config) + self._init_shared_buffers(config) + self._init_pool_scalars(config) + # Pre-create the side stream used by fused_topk_topp_renorm. Must + # happen before any CUDA graph capture — cudaStreamCreate is illegal + # inside capture, and verify() runs from the captured graph. + fused_topk_topp_prepare(config.device) + + def _init_dp_sampling(self, config: SamplingBackendConfig) -> None: + self._dp_tp_group = config.tp_group + self._dp_tp_size = ( + len(self._dp_tp_group) if self._dp_tp_group is not None else 1 + ) + self._dp_rank = 0 + self._dp_comm: DpSamplingComm | None = None + self._dp_comm_vocab_size = 0 + + if self._dp_tp_size <= 1: + self._dp_max_pad_bs = config.max_bs + self._dp_max_reqs_per_rank = config.max_bs + return + + self._dp_max_pad_bs = ( + (config.max_bs + self._dp_tp_size - 1) // self._dp_tp_size + ) * self._dp_tp_size + self._dp_max_reqs_per_rank = self._dp_max_pad_bs // self._dp_tp_size + + def configure_dp_sampling(self, runtime: DpSamplingRuntimeConfig) -> None: + if not runtime.enabled: + return + if ( + runtime.vocab_size is None + or runtime.max_bucket_bs is None + or runtime.topology is None + or runtime.device is None + ): + raise RuntimeError("enabled DP sampling runtime is incomplete") + topology = runtime.topology + if topology.tp_size != self._dp_tp_size: + raise RuntimeError( + f"DP sampling runtime tp_size={topology.tp_size} " + f"does not match backend tp_size={self._dp_tp_size}" + ) + if topology.tp_group != self._dp_tp_group: + raise RuntimeError("DP sampling runtime tp_group does not match backend") + if self._dp_tp_group is None: + raise RuntimeError("dp_sampling requires a tp_group") + self._dp_rank = topology.tp_rank + if runtime.max_bucket_bs > self._dp_max_pad_bs: + raise RuntimeError( + f"DP sampling max_bucket_bs={runtime.max_bucket_bs} exceeds " + f"backend max_pad_bs={self._dp_max_pad_bs}" + ) + if runtime.vocab_size % self._dp_tp_size != 0: + raise RuntimeError( + f"DP sampling vocab_size={runtime.vocab_size} must be divisible by " + f"tp_size={self._dp_tp_size}" + ) + self._init_dp_verify_buffers(runtime.device) + if runtime.vocab_size == self._dp_comm_vocab_size: + return + if self._dp_comm is not None and self._dp_comm.is_initialized: + raise RuntimeError("Cannot resize DP sampling comm after use") + self._dp_comm_vocab_size = runtime.vocab_size + self._dp_comm = DpSamplingComm( + tp_size=self._dp_tp_size, + rank=self._dp_rank, + group=self._dp_tp_group, + max_pad_bs=self._dp_max_pad_bs, + num_tokens_per_req=runtime.num_tokens_per_req, + vocab_size=runtime.vocab_size, + logits_dtype=None, + device=runtime.device, + ) + + def _init_dp_verify_buffers(self, device: torch.device | str) -> None: + if self._predict_local_buf is not None: + return + + max_n = self.config.max_draft_tokens_per_req + self._predict_local_buf = torch.zeros( + (self._dp_max_reqs_per_rank * max_n,), dtype=torch.int32, device=device + ) + self._accept_index_local_buf = torch.zeros( + (self._dp_max_reqs_per_rank * max_n,), dtype=torch.int32, device=device + ) + self._accept_length_local_buf = torch.zeros( + (self._dp_max_reqs_per_rank,), dtype=torch.int32, device=device + ) + + def _init_pool_scalars(self, config: SamplingBackendConfig) -> None: + # Capture warm-up reads row 0 with req_pool_indices zeroed, so row 0 + # must carry neutral-sampling values that can't produce nan/inf. + pool_rows = config.max_req_pool_size + 1 + + self._temperature_pool = torch.ones( + (pool_rows,), dtype=torch.float32, device=config.device + ) + self._top_k_pool = torch.ones( + (pool_rows,), dtype=torch.int32, device=config.device + ) + self._top_p_pool = torch.ones( + (pool_rows,), dtype=torch.float32, device=config.device + ) + self._seed_pool = torch.zeros( + (pool_rows,), dtype=torch.int64, device=config.device + ) + + # Per-slot CPU-side torch.Generators used to advance speculative + # coin buffers outside the CUDA graph. Seeded on flip from sp.seed. + # Slot 0 is pre-filled with _capture_gen so capture warm-up works + # without any real request having been registered. + # + # Retract-resume note: if a request is retracted and later takes a + # different pool slot on resume, _reset_slot re-seeds a fresh + # Generator from sp.seed. Sampling stays deterministic given the same + # seed, and flashinfer's Philox path (seed + seq_len offset) already + # gives per-step uniqueness independent of the torch.Generator. + self._cpu_generator_per_slot: list[torch.Generator | None] = [None] * pool_rows + self._cpu_generator_per_slot[0] = self._capture_gen + + def _reset_slot(self, pool_idx: int, sp: SamplingParams) -> None: + self._temperature_pool[pool_idx].fill_(float(sp.temperature)) + self._top_k_pool[pool_idx].fill_(int(sp.top_k)) + self._top_p_pool[pool_idx].fill_(float(sp.top_p)) + self._seed_pool[pool_idx].fill_(int(sp.seed)) + + cpu_gen = torch.Generator(device="cpu") + cpu_gen.manual_seed(int(sp.seed)) + self._cpu_generator_per_slot[pool_idx] = cpu_gen + + def _init_shared_buffers(self, config: SamplingBackendConfig) -> None: + + max_pad_bs = self._dp_max_pad_bs + max_n = config.max_draft_tokens_per_req + # Persistent coin buffers. Filled per-request in prepare() outside the + # CUDA graph so verify() only reads from them. + self._coins_buf = torch.zeros( + (max_pad_bs, max_n), + dtype=torch.float32, + device=config.device, + ) + self._final_coins_buf = torch.zeros( + (max_pad_bs,), dtype=torch.float32, device=config.device + ) + + # Stub generator used during CUDA-graph capture/warm-up (no requests yet). + self._capture_gen = torch.Generator(device=config.device) + self._capture_gen.manual_seed(config.random_seed) + + # Pre-allocated persistent buffers — no per-step alloc in the hot path. + self._ones_buf = torch.ones( + (max_pad_bs,), dtype=torch.int32, device=config.device + ) + # predict + accept_length share one packed backing store. + # Layout: [0, max_bs * max_n) is predict, [max_bs * max_n, total) + # is accept_length. + self._predict_max = max_pad_bs * max_n + self._output_pack_buf = torch.zeros( + (self._predict_max + max_pad_bs,), + dtype=torch.int32, + device=config.device, + ) + self._predict_buf = self._output_pack_buf[: self._predict_max] + self._accept_length_buf = self._output_pack_buf[self._predict_max :] + # Flat layout so [:bs * n].view(bs, n) is contiguous for any bs/n. + self._accept_index_buf = torch.zeros( + (max_pad_bs * max_n,), + dtype=torch.int32, + device=config.device, + ) + + self._predict_local_buf: torch.Tensor | None = None + self._accept_index_local_buf: torch.Tensor | None = None + self._accept_length_local_buf: torch.Tensor | None = None + + @torch.compile(dynamic=True, backend=get_compiler_backend()) + def _prepare_step_hook( + self, + num_tokens_per_req: int, + bs: int, + request_pool_indices: list[int] | None = None, + ) -> None: + """Refill persistent coin buffers outside the captured graph. + request_pool_indices=None is the capture/warm-up path — uses + _capture_gen for all rows. Otherwise reads per-slot generators + populated via _reset_slot.""" + if bs <= 0: + return + + n = min(num_tokens_per_req, self.config.max_draft_tokens_per_req) + lo = coin_eps(self._coins_buf.dtype) + if request_pool_indices is None: + self._coins_buf[:bs, :n].uniform_(lo, 1.0, generator=self._capture_gen) + self._final_coins_buf[:bs].uniform_(lo, 1.0, generator=self._capture_gen) + return + + cpu_coins = torch.empty((bs, n), dtype=torch.float32, pin_memory=True) + cpu_final = torch.empty((bs,), dtype=torch.float32, pin_memory=True) + + for i, pool_idx in enumerate(request_pool_indices): + gen = self._cpu_generator_per_slot[pool_idx] + if gen is None: + raise RuntimeError( + f"sampling slot {pool_idx} was not initialized before " + "coin-buffer refill" + ) + cpu_coins[i, :n].uniform_(lo, 1.0, generator=gen) + cpu_final[i].uniform_(lo, 1.0, generator=gen) + + self._coins_buf[:bs, :n].copy_(cpu_coins, non_blocking=True) + self._final_coins_buf[:bs].copy_(cpu_final, non_blocking=True) + + @nvtx_range("sampling:sample", color="yellow") + def sample( + self, + logits_output: LogitsProcessorOutput, + sampling_info: SamplingBatchInfo, + ) -> tuple[torch.Tensor, torch.Tensor]: + + logits = logits_output.next_token_logits + + # Grammar bitmask apply — captured inside the CUDA graph. Buffer is + # pre-bound by bind_grammar_mask_buf; non-grammar rows stay all-ones. + if sampling_info.vocab_mask is not None: + sampling_info.apply_vocab_mask( + logits=logits, vocab_mask=sampling_info.vocab_mask + ) + + if sampling_info.is_all_greedy: + + batch_next_token_ids = sampling_argmax(logits) + + else: + + temperatures, top_ks, top_ps, _, seeds, offsets = gather_and_expand_scalars( + sampling_info.req_pool_indices, + temperature=self._temperature_pool, + top_k=self._top_k_pool, + top_p=self._top_p_pool, + seed=self._seed_pool, + offsets=sampling_info.valid_cache_lengths, + enable_pdl=pdl_enabled(), + ) + + probs = softmax( + logits, + temperature=temperatures.view(-1, 1), + enable_pdl=pdl_enabled(), + ) + batch_next_token_ids = top_k_top_p_sampling_from_probs( + probs, + top_ks, + top_ps, + filter_apply_order="joint", + seed=seeds, + offset=offsets, + deterministic=True, + ) + + sampled = batch_next_token_ids.to(torch.int32) + + # TP-rank sync: rank 0 wins. + self.maybe_broadcast(sampled) + + if self.config.enable_output_logprobs: + logits_output.next_token_logprobs = gather_token_logprobs_torch( + logits, sampled + ) + + bs = logits.shape[0] + + return sampled, self._ones_buf[:bs] + + @nvtx_range("sampling:verify", color="yellow") + def verify( + self, + logits_output: LogitsProcessorOutput, + sampling_info: SamplingBatchInfo, + candidates: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + + bs = candidates.shape[0] + num_tokens_per_req = candidates.shape[1] + vocab_mask = sampling_info.vocab_mask + logits_layout_plan = getattr(logits_output, "logits_layout_plan", None) + dp_sampling = logits_layout_plan is not None + + if dp_sampling: + if self._dp_comm is None: + raise RuntimeError( + "dp_sampling requires tp_size > 1, a resolved tp_group, " + "and a configured DP comm" + ) + dp_comm = self._dp_comm + tp_size = self._dp_tp_size + rank = self._dp_rank + effective_bs = logits_layout_plan.effective_bs + pad_bs = logits_layout_plan.bucket_bs + if effective_bs != bs: + raise RuntimeError( + f"DP sampling effective_bs={effective_bs} must match " + f"candidate batch size {bs}" + ) + if ( + pad_bs < effective_bs + or pad_bs > self._dp_max_pad_bs + or pad_bs % tp_size != 0 + ): + raise RuntimeError( + f"invalid DP sampling pad_bs={pad_bs} for effective_bs={effective_bs}, " + f"max_pad_bs={self._dp_max_pad_bs}, tp_size={tp_size}" + ) + bs = pad_bs // tp_size + + # Shard by request so each request's draft chain stays on one rank. + shard = slice(rank * bs, (rank + 1) * bs) + if pad_bs > effective_bs: + candidates = torch.nn.functional.pad( + candidates, (0, 0, 0, pad_bs - effective_bs) + )[shard] + pool_indices = torch.nn.functional.pad( + sampling_info.req_pool_indices, (0, pad_bs - effective_bs) + )[shard] + else: + candidates = candidates[shard] + pool_indices = sampling_info.req_pool_indices[shard] + vocab_mask = slice_dp_vocab_mask( + vocab_mask, + full_bs=effective_bs, + pad_bs=pad_bs, + num_tokens_per_req=num_tokens_per_req, + shard=shard, + ) + coins = self._coins_buf[shard] + final_coins = self._final_coins_buf[shard] + if ( + self._predict_local_buf is None + or self._accept_index_local_buf is None + or self._accept_length_local_buf is None + ): + raise RuntimeError("DP sampling verify buffers are not initialized") + predict = self._predict_local_buf[: bs * num_tokens_per_req] + accept_index = ( + self._accept_index_local_buf[: bs * num_tokens_per_req] + .view(bs, num_tokens_per_req) + .fill_(-1) + ) + accept_length = self._accept_length_local_buf[:bs] + else: + pool_indices = sampling_info.req_pool_indices + coins = self._coins_buf + final_coins = self._final_coins_buf + predict = self._predict_buf[: bs * num_tokens_per_req] + accept_index = ( + self._accept_index_buf[: bs * num_tokens_per_req] + .view(bs, num_tokens_per_req) + .fill_(-1) + ) + accept_length = self._accept_length_buf[:bs] + + logits = logits_output.next_token_logits + if dp_sampling: + expected_rows = bs * num_tokens_per_req + if logits.shape[0] != expected_rows: + raise RuntimeError( + f"DP sampling logits rows {logits.shape[0]} != expected " + f"{expected_rows}" + ) + + # Per-draft-position grammar bitmask: buffer shape + # [bs * num_tokens_per_req, V/32] matches the flat target logits. + if vocab_mask is not None: + sampling_info.apply_vocab_mask( + logits=logits, + vocab_mask=vocab_mask, + ) + + if sampling_info.is_all_greedy: + + target_predict = sampling_argmax(logits).reshape(bs, num_tokens_per_req) + + verify_chain_greedy( + predicts=predict, + accept_index=accept_index, + accept_token_num=accept_length, + candidates=candidates, + target_predict=target_predict, + batch_size=bs, + num_draft_tokens=num_tokens_per_req, + enable_pdl=pdl_enabled(), + ) + + else: + + # Each request's N verified positions share one (temp, top_k, top_p) + # tuple; flat [bs*N] per-row knobs match the flat [bs*N, vocab] logits. + n = num_tokens_per_req + temperatures, top_ks, top_ps, _, _, _ = gather_and_expand_scalars( + pool_indices, + temperature=self._temperature_pool, + top_k=self._top_k_pool, + top_p=self._top_p_pool, + n=n, + enable_pdl=pdl_enabled(), + ) + + target_probs = softmax( + logits, + temperature=temperatures, + enable_pdl=pdl_enabled(), + ) + if _FUSED_TOPK_TOPP_AVAILABLE: + # Fused replacement for the back-to-back top_k_renorm_prob + + # top_p_renorm_prob(is_deterministic=True) pair. Sentinel + # K = 1<<30 in top_ks routes per-row through the radix top-p + # only path. + target_probs = fused_topk_topp_renorm( + target_probs, + top_ks, + top_ps, + enable_pdl=pdl_enabled(), + ) + else: + target_probs = top_k_renorm_prob(target_probs, top_ks) + target_probs = top_p_renorm_prob( + target_probs, top_ps, is_deterministic=True + ) + target_probs = target_probs.reshape(bs, n, -1) + + chain_speculative_sampling_target_only( + predicts=predict, + accept_index=accept_index, + accept_token_num=accept_length, + candidates=candidates, + uniform_samples=coins[:bs, :n], + uniform_samples_for_final_sampling=final_coins[:bs], + target_probs=target_probs, + draft_probs=None, + threshold_single=SPECULATIVE_ACCEPT_THRESHOLD_SINGLE, + threshold_acc=SPECULATIVE_ACCEPT_THRESHOLD_ACC, + deterministic=not dp_sampling, + enable_pdl=pdl_enabled(), + ) + + accept_length += 1 + logprobs_local = None + if self.config.enable_output_logprobs and dp_sampling: + # DP verify logits are still sharded by request at this point. + # Compute scalar logprobs for local predictions before gathering + # predictions to full-batch shape; the non-DP writer requires + # matching logits/token row counts. + logprobs_local = gather_token_logprobs_torch(logits, predict).view( + bs, num_tokens_per_req + ) + + if dp_sampling: + n = num_tokens_per_req + dp_comm.prepare_verify_outputs(logits_output.next_token_logits.dtype) + ( + predict_full, + accept_index_full, + accept_length_full, + ) = dp_comm.gather_verify_outputs( + predict_local=predict.view(bs, n), + accept_index_local=accept_index, + accept_length_local=accept_length, + pad_bs=pad_bs, + ) + predict = predict_full.view(-1)[: effective_bs * n] + accept_index = accept_index_full[:effective_bs] + accept_length = accept_length_full[:effective_bs] + if logprobs_local is not None: + logprobs_full = dp_comm.gather_verify_logprobs( + logprobs_local, + pad_bs=pad_bs, + ) + logits_output.next_token_logprobs = logprobs_full.view(-1)[ + : effective_bs * n + ] + # TP-rank sync: rank 0 wins on the full verify-output triple. + # Load-bearing: flashinfer top_k_renorm_prob has no is_deterministic + # knob and produces non-bit-identical results across ranks (sub-ulp + # FP accumulation order). + # PDL still uses rank-0 outputs to keep ranks aligned. Without PDL, + # fused top-k + top-p is bit-identical across ranks and does not need + # a broadcast. + elif pdl_enabled(): + self.maybe_broadcast(predict, accept_index, accept_length) + elif not _FUSED_TOPK_TOPP_AVAILABLE: + self.maybe_broadcast(predict, accept_index, accept_length) + + if self.config.enable_output_logprobs and not dp_sampling: + logits_output.next_token_logprobs = gather_token_logprobs_torch( + logits, predict + ) + + return predict, accept_length + + def get_packed_output_d2h( + self, + output_tokens: torch.Tensor, + output_lengths: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor] | None: + """One D2H of the packed predict+accept_length region. + + Only applies when both outputs alias into ``_output_pack_buf`` (the + verify() path). For ``sample()``, ``output_tokens`` is a fresh + argmax/top_k_top_p result and ``output_lengths`` is ``_ones_buf``, + neither of which lives in the pack. We fall back to two D2Hs. + """ + if ( + output_tokens.data_ptr() != self._output_pack_buf.data_ptr() + or output_lengths.data_ptr() != self._accept_length_buf.data_ptr() + ): + return None + n_t = output_tokens.numel() + n_l = output_lengths.numel() + # Copy the whole [0, predict_max + n_l). The gap [n_t, predict_max) + # is stale padding (max_bs * max_n vs. bs * n) — small enough that + # the saved launch beats the wasted bandwidth. + size = self._predict_max + n_l + cpu_pack = torch.empty(size, dtype=torch.int32, pin_memory=True) + cpu_pack.copy_(self._output_pack_buf[:size], non_blocking=True) + return ( + cpu_pack[:n_t].view(output_tokens.shape), + cpu_pack[self._predict_max : self._predict_max + n_l], + ) + + +register_backend("flashinfer", FlashInferSamplingBackend) diff --git a/python/tokenspeed/runtime/sampling/backends/flashinfer_full.py b/python/tokenspeed/runtime/sampling/backends/flashinfer_full.py new file mode 100644 index 0000000..f389277 --- /dev/null +++ b/python/tokenspeed/runtime/sampling/backends/flashinfer_full.py @@ -0,0 +1,495 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from tokenspeed_kernel.ops.sampling.cuda import ( + chain_speculative_sampling_target_only, + fused_topk_topp_renorm, +) +from tokenspeed_kernel.ops.sampling.flashinfer import ( + min_p_sampling_from_probs, + softmax, + top_k_renorm_prob, + top_p_renorm_prob, +) +from tokenspeed_kernel.ops.sampling.triton import ( + gather_and_expand_scalars, + min_p_renorm_prob, +) +from tokenspeed_kernel.torch_compile import get_compiler_backend + +from tokenspeed.runtime.sampling.backends.base import ( + SPECULATIVE_ACCEPT_THRESHOLD_ACC, + SPECULATIVE_ACCEPT_THRESHOLD_SINGLE, + SamplingBackendConfig, +) +from tokenspeed.runtime.sampling.backends.flashinfer import ( + _FUSED_TOPK_TOPP_AVAILABLE, + FlashInferSamplingBackend, +) +from tokenspeed.runtime.sampling.registry import register_backend +from tokenspeed.runtime.utils.nvtx import nvtx_range +from tokenspeed.runtime.utils.pdl import pdl_enabled + +if TYPE_CHECKING: + from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput + from tokenspeed.runtime.sampling.sampling_batch_info import SamplingBatchInfo + from tokenspeed.runtime.sampling.sampling_params import SamplingParams + + +class FlashInferFullSamplingBackend(FlashInferSamplingBackend): + """Superset of `flashinfer` adding min_p, frequency/presence/repetition + penalties, and per-token logit_bias, for both single-step sampling and + multi-step spec-decode verification. + + Stochastic path runs the 4-kernel sequence softmax(temperature) → + top_k_renorm → top_p_renorm → min_p_sampling, unconditionally (requests + with min_p == 0 are a no-op through min_p_sampling_from_probs) so the + captured CUDA graph matches the runtime flow. + + Layout: + * Per-pool-idx token counts (int32[max_req_pool_size, vocab]) — + accumulated after each sample/verify. Zeroed when a pool slot is + re-assigned to a new rid (see `on_pool_assignment`). + * Per-pool-idx logit bias (bf16[max_req_pool_size, vocab]) — zero by + default, scattered from SamplingParams.logit_bias on pool + assignment. Added to logits per step. + * Per-batch-row bf16 penalty scalars flowing through SamplingBatchInfo. + + sample() / verify() apply (in order, BEFORE temperature/softmax): + 1. repetition (multiplicative): logits = where(count>0, + where(logits>0, logits/rep, logits*rep), logits) + 2. frequency + presence (additive): + logits -= freq_pen * count + pres_pen * (count>0) + 3. logit_bias (additive): logits += logit_bias[req_pool_idx] + + Post-sample/verify, accumulate accepted tokens into counts. + + Out of scope in this iteration: min_new_tokens EOS mask, grammar vocab + mask. Both remain silently-ignored no-ops. + """ + + _SUPPORTS_DP_VERIFY = False + + def __init__(self, config: SamplingBackendConfig) -> None: + + super().__init__(config) + + if config.max_req_pool_size <= 0 or config.vocab_size <= 0: + + raise ValueError( + "FlashInferFullSamplingBackend requires max_req_pool_size > 0 and " + f"vocab_size > 0; got max_req_pool_size={config.max_req_pool_size}, " + f"vocab_size={config.vocab_size}" + ) + + # Valid pool indices run 0..max_req_pool_size inclusive. + pool_rows = config.max_req_pool_size + 1 + + self._counts = torch.zeros( + (pool_rows, config.vocab_size), + dtype=torch.int32, + device=config.device, + ) + + # bf16 is enough precision for typical client-supplied bias values + # (OpenAI caps |logit_bias| at 100). + self._logit_bias = torch.zeros( + (pool_rows, config.vocab_size), + dtype=torch.bfloat16, + device=config.device, + ) + + # Per-request penalty scalars + min_p. rep_pen starts at 1.0 + # (multiplicative identity); others at 0.0 (additive identity). + self._min_p_pool = torch.zeros( + (pool_rows,), dtype=torch.float32, device=config.device + ) + self._freq_pen_pool = torch.zeros( + (pool_rows,), dtype=torch.bfloat16, device=config.device + ) + self._pres_pen_pool = torch.zeros( + (pool_rows,), dtype=torch.bfloat16, device=config.device + ) + self._rep_pen_pool = torch.full( + (pool_rows,), 1.0, dtype=torch.bfloat16, device=config.device + ) + + # ------------------------------------------------------------------ + # Lifecycle hooks + # ------------------------------------------------------------------ + + def _reset_slot(self, pool_idx: int, sp: SamplingParams) -> None: + + # Scatter scalars inherited from the parent backend (temperature, top_k, + # top_p, seed). + super()._reset_slot(pool_idx, sp) + + # Penalty + min_p scalars. + self._min_p_pool[pool_idx].fill_(float(sp.min_p)) + self._freq_pen_pool[pool_idx].fill_(float(sp.frequency_penalty)) + self._pres_pen_pool[pool_idx].fill_(float(sp.presence_penalty)) + self._rep_pen_pool[pool_idx].fill_(float(sp.repetition_penalty)) + + # Zero the slot's count row (history from the previous occupant is + # no longer applicable). + self._counts[pool_idx].fill_(0) + + # Zero + scatter logit_bias for the new rid. Zeroing the whole row + # first rather than diffing because the previous occupant's bias + # keys are unknown here. + self._logit_bias[pool_idx].fill_(0.0) + + bias_map = getattr(sp, "logit_bias", None) if sp is not None else None + if bias_map: + vocab = self._logit_bias.shape[1] + raw_ids = [int(tid) for tid in bias_map.keys()] + assert all(0 <= tid < vocab for tid in raw_ids), ( + f"logit_bias contains out-of-vocab token id(s); " + f"vocab_size={vocab}, offending={[t for t in raw_ids if not 0 <= t < vocab]}" + ) + token_ids = torch.tensor( + raw_ids, + device=self._logit_bias.device, + dtype=torch.long, + ) + bias_values = torch.tensor( + list(bias_map.values()), + device=self._logit_bias.device, + dtype=torch.bfloat16, + ) + self._logit_bias[pool_idx, token_ids] = bias_values + + def reset_capture_state(self) -> None: + + # Warm-up iterations route all pool indices to row 0, which + # accumulates sampled tokens into _counts[0]. Zero it so the graph + # captures reads against a clean baseline. _logit_bias[0] is only + # written in on_pool_assignment, so it stays zero across warm-up. + self._counts[0].fill_(0) + + # ------------------------------------------------------------------ + # Penalty + bias application (shared by sample and verify) + # ------------------------------------------------------------------ + + @nvtx_range("sampling:penalties", color="yellow") + @torch.compile(dynamic=True, backend=get_compiler_backend()) + def _apply_penalties_and_bias( + self, + logits: torch.Tensor, + sampling_info: SamplingBatchInfo, + num_tokens_per_req: int = 1, + ) -> torch.Tensor: + """logits is [bs * num_tokens_per_req, V]. Penalty scalars are gathered + from the pool-indexed buffers. num_tokens_per_req > 1 is the spec-decode + verify() path where per-request scalars are repeat_interleave'd to + align with flat logits. + """ + + pool_idx = sampling_info.req_pool_indices + + if num_tokens_per_req > 1: + + pool_idx = torch.repeat_interleave(pool_idx, num_tokens_per_req, dim=0) + + counts = self._counts.index_select(0, pool_idx) # [bs*N, V] + active = counts > 0 + counts_f = counts.to(logits.dtype) + active_f = active.to(logits.dtype) + + # Gather per-request penalty scalars from the pool. [bs*N] → [bs*N, 1] + # for broadcast against [bs*N, V] logits. + rep = ( + self._rep_pen_pool.index_select(0, pool_idx).to(logits.dtype).unsqueeze(-1) + ) + freq = ( + self._freq_pen_pool.index_select(0, pool_idx).to(logits.dtype).unsqueeze(-1) + ) + presence = ( + self._pres_pen_pool.index_select(0, pool_idx).to(logits.dtype).unsqueeze(-1) + ) + + # 1. Repetition (multiplicative). scales is 1.0 where count==0, else + # rep_pen. Apply as logits/scales where logits>0, logits*scales else. + scales = torch.where(active, rep.expand_as(logits), torch.ones_like(logits)) + logits = torch.where(logits > 0, logits / scales, logits * scales) + + # 2. Frequency + presence (additive). Fused into a single subtract. + logits = logits - freq * counts_f - presence * active_f + + # 3. Per-token logit_bias (additive). Rows without a logit_bias are + # all-zero, so the add is a no-op for them. + logits = logits + self._logit_bias.index_select(0, pool_idx) + + return logits + + @nvtx_range("sampling:accum_counts", color="yellow") + @torch.compile(dynamic=True, backend=get_compiler_backend()) + def _accumulate_counts( + self, + pool_idx: torch.Tensor, + tokens: torch.Tensor, + weights: torch.Tensor, + ) -> None: + """Graph-safe in-place scatter: counts[pool_idx, tokens] += weights. + weights is int32; 0 masks invalid rows, 1 accumulates.""" + + self._counts.index_put_( + (pool_idx, tokens.long()), + weights.to(torch.int32), + accumulate=True, + ) + + # ------------------------------------------------------------------ + # Sample / verify + # ------------------------------------------------------------------ + + @nvtx_range("sampling:sample", color="yellow") + def sample( + self, + logits_output: LogitsProcessorOutput, + sampling_info: SamplingBatchInfo, + ) -> tuple[torch.Tensor, torch.Tensor]: + + logits = logits_output.next_token_logits.float() + + # Grammar bitmask apply — captured inside the CUDA graph. Buffer is + # pre-bound by bind_grammar_mask_buf; non-grammar rows stay all-ones. + # Applied before raw_logprobs capture so constrained logprobs reflect + # the grammar-masked distribution. + if sampling_info.vocab_mask is not None: + sampling_info.apply_vocab_mask( + logits=logits, vocab_mask=sampling_info.vocab_mask + ) + + # Raw-distribution logprobs (pre-penalty, pre-temperature) when the + # server flag is on. Gather is done after we know the sampled id. + raw_logprobs = ( + torch.log_softmax(logits, dim=-1) + if self.config.enable_output_logprobs + else None + ) + + logits = self._apply_penalties_and_bias(logits, sampling_info) + + temperatures, top_ks, top_ps, min_ps, seeds, offsets = ( + gather_and_expand_scalars( + sampling_info.req_pool_indices, + temperature=self._temperature_pool, + top_k=self._top_k_pool, + top_p=self._top_p_pool, + min_p=self._min_p_pool, + seed=self._seed_pool, + offsets=sampling_info.valid_cache_lengths, + enable_pdl=pdl_enabled(), + ) + ) + + probs = softmax( + logits, temperature=temperatures.view(-1, 1), enable_pdl=pdl_enabled() + ) + + if _FUSED_TOPK_TOPP_AVAILABLE: + # Fused replacement for the back-to-back top_k_renorm_prob + + # top_p_renorm_prob(is_deterministic=True) pair. Sentinel + # K = 1<<30 in top_ks routes per-row through the radix top-p + # only path. + probs = fused_topk_topp_renorm( + probs, + top_ks, + top_ps, + enable_pdl=pdl_enabled(), + ) + else: + probs = top_k_renorm_prob(probs, top_ks) + probs = top_p_renorm_prob(probs, top_ps, is_deterministic=True) + + batch_next_token_ids = min_p_sampling_from_probs( + probs, + min_ps, + seed=seeds, + offset=offsets, + deterministic=True, + ) + + sampled = batch_next_token_ids.to(torch.int32) + + # TP-rank sync BEFORE _accumulate_counts so per-rank counts stay aligned. + # For fused top-k + top-p, the results are bit-identical across ranks. + # So we don't need to broadcast the results. + if not _FUSED_TOPK_TOPP_AVAILABLE: + self.maybe_broadcast(sampled) + + if raw_logprobs is not None: + + logits_output.next_token_logprobs = raw_logprobs.gather( + -1, sampled.unsqueeze(-1) + ).squeeze(-1) + + # Accumulate sampled tokens into counts (greedy path accumulates too + # so mixed later batches see the correct history). + self._accumulate_counts( + sampling_info.req_pool_indices, + sampled, + torch.ones_like(sampled, dtype=torch.int32), + ) + + bs = logits.shape[0] + + return sampled, self._ones_buf[:bs] + + @nvtx_range("sampling:verify", color="yellow") + def verify( + self, + logits_output: LogitsProcessorOutput, + sampling_info: SamplingBatchInfo, + candidates: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + + bs = candidates.shape[0] + num_tokens_per_req = candidates.shape[1] + + predict = self._predict_buf[: bs * num_tokens_per_req] + accept_index = ( + self._accept_index_buf[: bs * num_tokens_per_req] + .view(bs, num_tokens_per_req) + .fill_(-1) + ) + accept_length = self._accept_length_buf[:bs] + + logits = logits_output.next_token_logits.float() + + # Per-draft-position grammar bitmask: buffer shape + # [bs * num_tokens_per_req, V/32] matches the flat target logits. + # Applied before raw_logprobs capture so constrained logprobs reflect + # the grammar-masked distribution. + if sampling_info.vocab_mask is not None: + sampling_info.apply_vocab_mask( + logits=logits, + vocab_mask=sampling_info.vocab_mask, + ) + + # Raw (pre-penalty) logprobs captured before penalty application to + # match sample()'s semantics. + raw_logprobs = ( + torch.log_softmax(logits, dim=-1) + if self.config.enable_output_logprobs + else None + ) + + logits = self._apply_penalties_and_bias( + logits, + sampling_info, + num_tokens_per_req=num_tokens_per_req, + ) + + temperatures, top_ks, top_ps, min_ps, _, _ = gather_and_expand_scalars( + sampling_info.req_pool_indices, + temperature=self._temperature_pool, + top_k=self._top_k_pool, + top_p=self._top_p_pool, + min_p=self._min_p_pool, + n=num_tokens_per_req, + enable_pdl=pdl_enabled(), + ) + + target_probs = softmax( + logits, temperature=temperatures.view(-1, 1), enable_pdl=pdl_enabled() + ) + if _FUSED_TOPK_TOPP_AVAILABLE: + # Fused replacement for the back-to-back top_k_renorm_prob + + # top_p_renorm_prob(is_deterministic=True) pair. Sentinel + # K = 1<<30 in top_ks routes per-row through the radix top-p + # only path. + target_probs = fused_topk_topp_renorm( + target_probs, + top_ks, + top_ps, + enable_pdl=pdl_enabled(), + ) + else: + target_probs = top_k_renorm_prob(target_probs, top_ks) + target_probs = top_p_renorm_prob( + target_probs, top_ps, is_deterministic=True + ) + + target_probs = min_p_renorm_prob(target_probs, min_ps, enable_pdl=pdl_enabled()) + + target_probs = target_probs.reshape(bs, num_tokens_per_req, -1) + + coins = self._coins_buf[:bs, :num_tokens_per_req] + coins_for_final_sampling = self._final_coins_buf[:bs] + + chain_speculative_sampling_target_only( + predicts=predict, + accept_index=accept_index, + accept_token_num=accept_length, + candidates=candidates.to(torch.int32), + uniform_samples=coins, + uniform_samples_for_final_sampling=coins_for_final_sampling, + target_probs=target_probs, + draft_probs=None, + threshold_single=SPECULATIVE_ACCEPT_THRESHOLD_SINGLE, + threshold_acc=SPECULATIVE_ACCEPT_THRESHOLD_ACC, + deterministic=True, + enable_pdl=pdl_enabled(), + ) + + accept_length += 1 + + # TP-rank sync BEFORE _accumulate_counts so per-rank counts stay aligned. + # For fused top-k + top-p, the results are bit-identical across ranks. + # So we don't need to broadcast the results. + if not _FUSED_TOPK_TOPP_AVAILABLE: + self.maybe_broadcast(predict, accept_index, accept_length) + + # Accumulate accepted tokens into counts. accept_index is [bs, N] + # with -1 in unused slots; clamp to a safe index and mask with a + # weight of 0 so invalid slots are no-ops. + valid = accept_index >= 0 # [bs, N] + safe_positions = accept_index.clamp(min=0).long() # [bs, N] + accepted_tokens = predict.long().gather(0, safe_positions.view(-1)) + + pool_idx_expanded = ( + sampling_info.req_pool_indices.unsqueeze(-1) + .expand(-1, num_tokens_per_req) + .reshape(-1) + ) + + self._accumulate_counts( + pool_idx_expanded, + accepted_tokens, + valid.reshape(-1).to(torch.int32), + ) + + if raw_logprobs is not None: + + logits_output.next_token_logprobs = raw_logprobs.gather( + -1, predict.unsqueeze(-1) + ).squeeze(-1) + + return predict, accept_length + + +register_backend("flashinfer_full", FlashInferFullSamplingBackend) diff --git a/python/tokenspeed/runtime/sampling/backends/greedy.py b/python/tokenspeed/runtime/sampling/backends/greedy.py new file mode 100644 index 0000000..97dc4b2 --- /dev/null +++ b/python/tokenspeed/runtime/sampling/backends/greedy.py @@ -0,0 +1,253 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from tokenspeed_kernel.ops.sampling import argmax as sampling_argmax +from tokenspeed_kernel.ops.sampling.cuda import ( + verify_chain_greedy as _verify_chain_greedy_cuda, +) +from tokenspeed_kernel.registry import error_fn + +from tokenspeed.runtime.sampling.backends.base import ( + SamplingBackend, + SamplingBackendConfig, +) +from tokenspeed.runtime.sampling.registry import register_backend +from tokenspeed.runtime.sampling.utils import gather_token_logprobs_torch +from tokenspeed.runtime.utils.nvtx import nvtx_range +from tokenspeed.runtime.utils.pdl import pdl_enabled + +if TYPE_CHECKING: + + from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput + from tokenspeed.runtime.sampling.sampling_batch_info import SamplingBatchInfo + + +def _verify_chain_greedy_torch( + predicts: torch.Tensor, # [bs * N] int32, in/out + accept_index: torch.Tensor, # [bs, N] int32, in/out (-1-filled on entry) + accept_token_num: torch.Tensor, # [bs] int32, out + candidates: torch.Tensor, # [bs, N] int32 + target_predict: torch.Tensor, # [bs, N] int64 (argmax output) + batch_size: int, + num_draft_tokens: int, +) -> None: + """Pure-torch equivalent of tokenspeed_kernel.verify_chain_greedy. + + Used on non-CUDA devices and when the CUDA kernel is unavailable. + """ + + bs = batch_size + n = num_draft_tokens + + # For i in 1..n-1: candidates[b, i] accepted iff it equals target_predict[b, i-1]. + # Accepted prefix length per row = longest-leading-1s of the match array. + match = candidates[:, 1:] == target_predict[:, :-1].to( + candidates.dtype + ) # [bs, n-1] + leading = torch.cumprod(match.to(torch.int32), dim=1) # [bs, n-1] + num_accepted = leading.sum(dim=1).to(torch.int32) # [bs] + + # Fill all of `predicts` with target_predict; slots outside the accepted + # prefix are harmless because accept_index keeps them at -1 and callers + # mask on that. Matches the CUDA kernel's observable state. + predicts.copy_(target_predict.reshape(-1).to(torch.int32)) + + device = candidates.device + pos = torch.arange(n, device=device).unsqueeze(0) # [1, n] + batch_off = torch.arange(bs, device=device).unsqueeze(1) * n # [bs, 1] + flat_idx = (batch_off + pos).to(torch.int32) # [bs, n] + valid = pos <= num_accepted.unsqueeze(1) # [bs, n] + accept_index.copy_(torch.where(valid, flat_idx, torch.full_like(accept_index, -1))) + accept_token_num.copy_(num_accepted) + + +def _verify_chain_greedy( + predicts: torch.Tensor, + accept_index: torch.Tensor, + accept_token_num: torch.Tensor, + candidates: torch.Tensor, + target_predict: torch.Tensor, + batch_size: int, + num_draft_tokens: int, + enable_pdl: bool = False, +) -> None: + + # Prefer the CUDA kernel when available AND the tensors are on CUDA. + if _verify_chain_greedy_cuda is not error_fn and candidates.is_cuda: + + _verify_chain_greedy_cuda( + predicts=predicts, + accept_index=accept_index, + accept_token_num=accept_token_num, + candidates=candidates, + target_predict=target_predict, + batch_size=batch_size, + num_draft_tokens=num_draft_tokens, + enable_pdl=enable_pdl, + ) + return + + _verify_chain_greedy_torch( + predicts=predicts, + accept_index=accept_index, + accept_token_num=accept_token_num, + candidates=candidates, + target_predict=target_predict, + batch_size=batch_size, + num_draft_tokens=num_draft_tokens, + ) + + +class GreedySamplingBackend(SamplingBackend): + """Greedy-only backend: argmax for single-step, chain-greedy verify for + multi-step verification. No flashinfer / min_p / penalty machinery, no + coin buffers. Verify uses the fused CUDA kernel when available; falls + back to a pure-torch implementation otherwise (CPU, ROCm, etc.). + + sampling_info is ignored for single-step (always argmax). verify() also + treats every request as greedy — stochastic verification is not + supported. Intended as the default backend and as a fallback when + flashinfer is unavailable.""" + + def __init__(self, config: SamplingBackendConfig) -> None: + + super().__init__(config) + + self._ones_buf = torch.ones( + (config.max_bs,), dtype=torch.int32, device=config.device + ) + # Pre-allocated int32 buffer for ``sample``'s argmax output: lets the + # cute_dsl kernel write int32 token ids directly, skipping the + # ``.to(torch.int32)`` cast and its elementwise launch in the + # CUDA-graph-captured hot path. + self._sample_token_buf = torch.empty( + (config.max_bs,), dtype=torch.int32, device=config.device + ) + self._predict_buf = torch.zeros( + (config.max_bs * config.max_draft_tokens_per_req,), + dtype=torch.int32, + device=config.device, + ) + # Flat layout so [:bs * n].view(bs, n) is contiguous for any bs/n + # (required by maybe_broadcast / NCCL). + self._accept_index_buf = torch.zeros( + (config.max_bs * config.max_draft_tokens_per_req,), + dtype=torch.int32, + device=config.device, + ) + self._accept_length_buf = torch.zeros( + (config.max_bs,), dtype=torch.int32, device=config.device + ) + + @nvtx_range("sampling:sample", color="yellow") + def sample( + self, + logits_output: LogitsProcessorOutput, + sampling_info: SamplingBatchInfo, + ) -> tuple[torch.Tensor, torch.Tensor]: + + logits = logits_output.next_token_logits + # Grammar bitmask apply — captured inside the CUDA graph. Buffer is + # pre-bound by bind_grammar_mask_buf; non-grammar rows stay all-ones + # so apply is a no-op. + if sampling_info.vocab_mask is not None: + sampling_info.apply_vocab_mask( + logits=logits, vocab_mask=sampling_info.vocab_mask + ) + bs = logits.shape[0] + tokens = sampling_argmax(logits, out=self._sample_token_buf[:bs]) + + # TP-rank sync (rank 0 wins), mirrors FlashInferSamplingBackend.sample. + # All-gathered logits are not bit-identical across ranks, so per-rank + # argmax can diverge; an unsynced token id desyncs batch composition and + # deadlocks a downstream model all-reduce. + self.maybe_broadcast(tokens) + + if self.config.enable_output_logprobs: + logits_output.next_token_logprobs = gather_token_logprobs_torch( + logits, tokens + ) + + return tokens, self._ones_buf[:bs] + + @nvtx_range("sampling:verify", color="yellow") + def verify( + self, + logits_output: LogitsProcessorOutput, + sampling_info: SamplingBatchInfo, + candidates: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + + bs = candidates.shape[0] + num_tokens_per_req = candidates.shape[1] + + predict = self._predict_buf[: bs * num_tokens_per_req] + accept_index = ( + self._accept_index_buf[: bs * num_tokens_per_req] + .view(bs, num_tokens_per_req) + .fill_(-1) + ) + accept_length = self._accept_length_buf[:bs] + + logits = logits_output.next_token_logits + + # Per-draft-position grammar bitmask: buffer shape + # [bs * num_tokens_per_req, V/32] matches the flat target logits. + if sampling_info.vocab_mask is not None: + sampling_info.apply_vocab_mask( + logits=logits, + vocab_mask=sampling_info.vocab_mask, + ) + target_predict = sampling_argmax(logits).reshape(bs, num_tokens_per_req) + + _verify_chain_greedy( + predicts=predict, + accept_index=accept_index, + accept_token_num=accept_length, + candidates=candidates.to(torch.int32), + target_predict=target_predict, + batch_size=bs, + num_draft_tokens=num_tokens_per_req, + enable_pdl=pdl_enabled(), + ) + + accept_length += 1 + + # TP-rank sync on the full verify-output triple, mirrors + # FlashInferSamplingBackend.verify. Per-rank argmax / accept-length + # divergence (logits not bit-identical across ranks) desyncs batch + # composition and deadlocks the model all-reduce. Buffers are laid out + # flat so these views are NCCL-contiguous. + self.maybe_broadcast(predict, accept_index, accept_length) + + if self.config.enable_output_logprobs: + logits_output.next_token_logprobs = gather_token_logprobs_torch( + logits, predict + ) + + return predict, accept_length + + +register_backend("greedy", GreedySamplingBackend) diff --git a/python/tokenspeed/runtime/sampling/backends/triton.py b/python/tokenspeed/runtime/sampling/backends/triton.py new file mode 100644 index 0000000..47698ff --- /dev/null +++ b/python/tokenspeed/runtime/sampling/backends/triton.py @@ -0,0 +1,707 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +# TokenSpeed keeps its request-pool state and backend boundary. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from tokenspeed_kernel.ops.sampling.cute_dsl import argmax as cute_argmax +from tokenspeed_kernel.ops.sampling.triton import ( + _QRITA_PERCENTILE_TO_STD_TABLE, + gumbel_sample_from_pools, + gumbel_sample_from_pools_compact, + gumbel_sample_from_pools_generic, + gumbel_sample_top_k_top_p_from_pools, + gumbel_sample_top_k_top_p_qrita_from_pools, + gumbel_sample_top_p_parallel_from_pools, + selected_token_logprobs, + verify_chain_target_sampled, +) + +from tokenspeed.runtime.sampling.backends.base import ( + CUDA_GRAPH_VARIANT_DEFAULT, + SamplingBackend, + SamplingBackendConfig, +) +from tokenspeed.runtime.sampling.registry import register_backend +from tokenspeed.runtime.sampling.sampling_params import _SAMPLING_EPS, _TOP_K_DISABLED +from tokenspeed.runtime.sampling.utils import nan_guard_logits +from tokenspeed.runtime.utils.nvtx import nvtx_range +from tokenspeed.runtime.utils.pdl import pdl_enabled + +if TYPE_CHECKING: + from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput + from tokenspeed.runtime.sampling.sampling_batch_info import SamplingBatchInfo + from tokenspeed.runtime.sampling.sampling_params import SamplingParams + + +_GUMBEL_BLOCK_SIZE = 1024 +_COMPACT_GUMBEL_BLOCK_SIZE = 4096 +_COMPACT_GUMBEL_VOCAB_MAX = 32768 +_TOP_K_TOP_P_SMALL_BLOCK_SIZE = 1024 +_TOP_K_TOP_P_TUNED_VOCAB_MAX = 32768 +_TOP_K_TOP_P_PAD = 128 +_TOP_P_PARALLEL_SAMPLE_BLOCK_SIZE = 1024 +_TOP_P_PARALLEL_SAMPLE_ATTEMPTS = 3 +_TOP_P_PARALLEL_VERIFY_ATTEMPTS = 4 +_TOP_P_PARALLEL_MAX_ATTEMPTS = max( + _TOP_P_PARALLEL_SAMPLE_ATTEMPTS, _TOP_P_PARALLEL_VERIFY_ATTEMPTS +) +_QRITA_VERIFY_MIN_ROWS = 128 + +_SAMPLE_ROUTE_GUMBEL_GENERIC = 0 +_SAMPLE_ROUTE_GUMBEL_NO_FILTER = 1 +_SAMPLE_ROUTE_GUMBEL_TOP_K = 2 +_SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P = 3 +_SAMPLE_ROUTE_GUMBEL_TOP_P = 4 +CUDA_GRAPH_VARIANT_TRITON_NO_FILTER = "triton_no_filter" +CUDA_GRAPH_VARIANT_TRITON_TOP_K = "triton_top_k" +CUDA_GRAPH_VARIANT_TRITON_TOP_K_TOP_P = "triton_top_k_top_p" +CUDA_GRAPH_VARIANT_TRITON_TOP_P = "triton_top_p" +CUDA_GRAPH_VARIANT_TRITON_VERIFY_NO_FILTER = "triton_verify_no_filter" + +_CUDA_GRAPH_VARIANT_SAMPLE_ROUTES = { + CUDA_GRAPH_VARIANT_TRITON_NO_FILTER: _SAMPLE_ROUTE_GUMBEL_NO_FILTER, + CUDA_GRAPH_VARIANT_TRITON_TOP_K: _SAMPLE_ROUTE_GUMBEL_TOP_K, + CUDA_GRAPH_VARIANT_TRITON_TOP_K_TOP_P: _SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P, + CUDA_GRAPH_VARIANT_TRITON_TOP_P: _SAMPLE_ROUTE_GUMBEL_TOP_P, + CUDA_GRAPH_VARIANT_TRITON_VERIFY_NO_FILTER: _SAMPLE_ROUTE_GUMBEL_NO_FILTER, +} + + +class TritonSamplingBackend(SamplingBackend): + """TokenSpeed pool-state backend using Triton Gumbel-Max kernels.""" + + _HAS_POOL_STATE = True + + def __init__(self, config: SamplingBackendConfig) -> None: + super().__init__(config) + self._init_triton_pool_state(config) + self._init_triton_buffers(config) + self._sample_route = _SAMPLE_ROUTE_GUMBEL_GENERIC + self._top_k_top_p_pad = _TOP_K_TOP_P_PAD + + def _init_triton_pool_state(self, config: SamplingBackendConfig) -> None: + pool_rows = config.max_req_pool_size + 1 + self._temperature_pool = torch.ones( + (pool_rows,), dtype=torch.float32, device=config.device + ) + self._top_k_pool = torch.ones( + (pool_rows,), dtype=torch.int32, device=config.device + ) + self._top_p_pool = torch.ones( + (pool_rows,), dtype=torch.float32, device=config.device + ) + self._seed_pool = torch.zeros( + (pool_rows,), dtype=torch.int64, device=config.device + ) + + self._ones_buf = torch.ones( + (config.max_bs,), dtype=torch.int32, device=config.device + ) + self._predict_buf = torch.zeros( + (config.max_bs * config.max_draft_tokens_per_req,), + dtype=torch.int32, + device=config.device, + ) + # Flat layout so [:bs * n].view(bs, n) is contiguous for any bs/n + # (required by maybe_broadcast / NCCL). + self._accept_index_buf = torch.zeros( + (config.max_bs * config.max_draft_tokens_per_req,), + dtype=torch.int32, + device=config.device, + ) + self._accept_length_buf = torch.zeros( + (config.max_bs,), dtype=torch.int32, device=config.device + ) + + def _init_triton_buffers(self, config: SamplingBackendConfig) -> None: + pool_rows = config.max_req_pool_size + 1 + self._zero_offsets_pool = torch.zeros( + (pool_rows,), dtype=torch.int64, device=config.device + ) + + vocab_size = max(int(config.vocab_size), 1) + gumbel_blocks = (vocab_size + _GUMBEL_BLOCK_SIZE - 1) // _GUMBEL_BLOCK_SIZE + self._gumbel_local_ids = torch.empty( + (config.max_bs, gumbel_blocks), + dtype=torch.int32, + device=config.device, + ) + self._gumbel_local_scores = torch.empty( + (config.max_bs, gumbel_blocks), + dtype=torch.float32, + device=config.device, + ) + self._gumbel_out = torch.empty( + (config.max_bs,), dtype=torch.int32, device=config.device + ) + self._req_pool_indices_i32 = torch.empty( + (config.max_bs,), dtype=torch.int32, device=config.device + ) + self._gumbel_verify_out = torch.empty( + (config.max_bs * config.max_draft_tokens_per_req,), + dtype=torch.int32, + device=config.device, + ) + self._gumbel_verify_local_ids = torch.empty( + (config.max_bs * config.max_draft_tokens_per_req, gumbel_blocks), + dtype=torch.int32, + device=config.device, + ) + self._gumbel_verify_local_scores = torch.empty( + (config.max_bs * config.max_draft_tokens_per_req, gumbel_blocks), + dtype=torch.float32, + device=config.device, + ) + + topk_blocks = (vocab_size + _TOP_K_TOP_P_SMALL_BLOCK_SIZE - 1) // ( + _TOP_K_TOP_P_SMALL_BLOCK_SIZE + ) + topk_candidates = topk_blocks * _TOP_K_TOP_P_PAD + self._topk_candidate_ids = torch.empty( + (config.max_bs, topk_candidates), + dtype=torch.int32, + device=config.device, + ) + self._topk_candidate_logits = torch.empty( + (config.max_bs, topk_candidates), + dtype=torch.float32, + device=config.device, + ) + self._topk_verify_candidate_ids = torch.empty( + (config.max_bs * config.max_draft_tokens_per_req, topk_candidates), + dtype=torch.int32, + device=config.device, + ) + self._topk_verify_candidate_logits = torch.empty( + (config.max_bs * config.max_draft_tokens_per_req, topk_candidates), + dtype=torch.float32, + device=config.device, + ) + max_verify_rows = max(config.max_bs * config.max_draft_tokens_per_req, 1) + num_sms = torch.cuda.get_device_properties(config.device).multi_processor_count + self._qrita_verify_num_programs = min(num_sms, max_verify_rows) + self._qrita_verify_buffer = torch.empty( + (self._qrita_verify_num_programs, vocab_size), + dtype=torch.float32, + device=config.device, + ) + self._qrita_percentile_to_std_table = torch.tensor( + _QRITA_PERCENTILE_TO_STD_TABLE, + dtype=torch.float32, + device=config.device, + ) + + top_p_blocks = (vocab_size + _TOP_P_PARALLEL_SAMPLE_BLOCK_SIZE - 1) // ( + _TOP_P_PARALLEL_SAMPLE_BLOCK_SIZE + ) + top_p_rows = max(config.max_bs * config.max_draft_tokens_per_req, 1) + self._top_p_local_max = torch.empty( + (top_p_rows, top_p_blocks), dtype=torch.float32, device=config.device + ) + self._top_p_local_sum = torch.empty( + (top_p_rows, top_p_blocks), dtype=torch.float32, device=config.device + ) + self._top_p_local_argmax = torch.empty( + (top_p_rows, top_p_blocks), dtype=torch.int32, device=config.device + ) + self._top_p_local_scores = torch.empty( + (top_p_rows, top_p_blocks, _TOP_P_PARALLEL_MAX_ATTEMPTS), + dtype=torch.float32, + device=config.device, + ) + self._top_p_local_logits = torch.empty( + (top_p_rows, top_p_blocks, _TOP_P_PARALLEL_MAX_ATTEMPTS), + dtype=torch.float32, + device=config.device, + ) + self._top_p_local_ids = torch.empty( + (top_p_rows, top_p_blocks, _TOP_P_PARALLEL_MAX_ATTEMPTS), + dtype=torch.int32, + device=config.device, + ) + self._top_p_row_max = torch.empty( + (top_p_rows,), dtype=torch.float32, device=config.device + ) + self._top_p_row_total = torch.empty( + (top_p_rows,), dtype=torch.float32, device=config.device + ) + self._top_p_row_argmax = torch.empty( + (top_p_rows,), dtype=torch.int32, device=config.device + ) + self._top_p_row_candidate_logits = torch.empty( + (top_p_rows, _TOP_P_PARALLEL_MAX_ATTEMPTS), + dtype=torch.float32, + device=config.device, + ) + self._top_p_row_candidate_ids = torch.empty( + (top_p_rows, _TOP_P_PARALLEL_MAX_ATTEMPTS), + dtype=torch.int32, + device=config.device, + ) + self._top_p_accepted = torch.empty( + (top_p_rows,), dtype=torch.int32, device=config.device + ) + self._selected_logprob_out = torch.empty( + (top_p_rows,), dtype=torch.float32, device=config.device + ) + + def _req_pool_indices_for_kernels( + self, req_pool_indices: torch.Tensor, rows: int + ) -> torch.Tensor: + req_pool_indices = req_pool_indices[:rows] + if req_pool_indices.dtype == torch.int32: + return req_pool_indices + if req_pool_indices.dtype != torch.int64: + raise ValueError( + "Triton sampling requires int32/int64 req_pool_indices, " + f"got {req_pool_indices.dtype}" + ) + out = self._req_pool_indices_i32[:rows] + out.copy_(req_pool_indices, non_blocking=True) + return out + + def _write_logprob_outputs( + self, + logits_output: LogitsProcessorOutput, + logits: torch.Tensor, + sampled: torch.Tensor, + ) -> None: + if not self.config.enable_output_logprobs: + return + + rows = logits.shape[0] + selected_out = self._selected_logprob_out[:rows] + logits_output.next_token_logprobs = selected_token_logprobs( + logits, sampled, selected_out + ) + + @staticmethod + def _select_sample_route( + sampling_params_list: list[SamplingParams], + ) -> int: + if len(sampling_params_list) == 0: + return _SAMPLE_ROUTE_GUMBEL_GENERIC + + top_ks = [int(sp.top_k) for sp in sampling_params_list] + top_ps = [float(sp.top_p) for sp in sampling_params_list] + all_top_p_one = all(abs(p - 1.0) <= _SAMPLING_EPS for p in top_ps) + all_top_k_disabled = all(k == _TOP_K_DISABLED for k in top_ks) + all_top_k_finite = all(k != _TOP_K_DISABLED for k in top_ks) + + if all_top_k_disabled and all_top_p_one: + return _SAMPLE_ROUTE_GUMBEL_NO_FILTER + if all_top_k_disabled: + return _SAMPLE_ROUTE_GUMBEL_TOP_P + if all_top_k_finite: + if all_top_p_one: + return _SAMPLE_ROUTE_GUMBEL_TOP_K + return _SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P + return _SAMPLE_ROUTE_GUMBEL_GENERIC + + def _reset_slot(self, pool_idx: int, sp: SamplingParams) -> None: + self._temperature_pool[pool_idx].fill_(float(sp.temperature)) + self._top_k_pool[pool_idx].fill_(int(sp.top_k)) + self._top_p_pool[pool_idx].fill_(float(sp.top_p)) + self._seed_pool[pool_idx].fill_(int(sp.seed)) + + def prepare_step( + self, + request_ids: list[str], + request_pool_indices: list[int], + sampling_params_list: list[SamplingParams], + num_tokens_per_req: int = 1, + ) -> None: + SamplingBackend.prepare_step( + self, + request_ids=request_ids, + request_pool_indices=request_pool_indices, + sampling_params_list=sampling_params_list, + num_tokens_per_req=num_tokens_per_req, + ) + self._sample_route = self._select_sample_route(sampling_params_list) + self._top_k_top_p_pad = self._select_top_k_top_p_pad(sampling_params_list) + + @staticmethod + def _select_top_k_top_p_pad(sampling_params_list: list[SamplingParams]) -> int: + finite_top_ks = [ + int(sp.top_k) + for sp in sampling_params_list + if int(sp.top_k) != _TOP_K_DISABLED + ] + if finite_top_ks and max(finite_top_ks) <= 64: + return 64 + return _TOP_K_TOP_P_PAD + + def _use_qrita_verify_top_k_route(self, rows: int, vocab_size: int) -> bool: + return ( + self._sample_route + in (_SAMPLE_ROUTE_GUMBEL_TOP_K, _SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P) + and rows >= _QRITA_VERIFY_MIN_ROWS + and vocab_size >= _TOP_K_TOP_P_TUNED_VOCAB_MAX + and ( + vocab_size > _TOP_K_TOP_P_TUNED_VOCAB_MAX or self._top_k_top_p_pad > 64 + ) + ) + + def prepare_capture(self, bs: int, num_tokens_per_req: int = 1) -> None: + self._sample_route = _SAMPLE_ROUTE_GUMBEL_GENERIC + self._top_k_top_p_pad = _TOP_K_TOP_P_PAD + SamplingBackend.prepare_capture( + self, bs=bs, num_tokens_per_req=num_tokens_per_req + ) + + def cuda_graph_capture_variants(self, num_tokens_per_req: int) -> tuple[str, ...]: + variants = ( + CUDA_GRAPH_VARIANT_DEFAULT, + CUDA_GRAPH_VARIANT_TRITON_NO_FILTER, + CUDA_GRAPH_VARIANT_TRITON_TOP_P, + CUDA_GRAPH_VARIANT_TRITON_TOP_K, + CUDA_GRAPH_VARIANT_TRITON_TOP_K_TOP_P, + ) + if num_tokens_per_req <= 1: + return variants + return (*variants, CUDA_GRAPH_VARIANT_TRITON_VERIFY_NO_FILTER) + + def prepare_capture_variant( + self, + bs: int, + num_tokens_per_req: int, + variant: str, + ) -> None: + sample_route = _CUDA_GRAPH_VARIANT_SAMPLE_ROUTES.get(variant) + if sample_route is not None: + self._sample_route = sample_route + if sample_route in ( + _SAMPLE_ROUTE_GUMBEL_TOP_K, + _SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P, + ): + self._top_k_top_p_pad = _TOP_K_TOP_P_PAD + SamplingBackend.prepare_capture( + self, + bs=bs, + num_tokens_per_req=num_tokens_per_req, + ) + return + if variant == CUDA_GRAPH_VARIANT_DEFAULT: + self.prepare_capture(bs=bs, num_tokens_per_req=num_tokens_per_req) + return + raise ValueError(f"Unsupported CUDA graph variant: {variant}") + + def cuda_graph_replay_variant(self, num_tokens_per_req: int) -> str: + if self._sample_route == _SAMPLE_ROUTE_GUMBEL_NO_FILTER: + if num_tokens_per_req > 1: + return CUDA_GRAPH_VARIANT_TRITON_VERIFY_NO_FILTER + return CUDA_GRAPH_VARIANT_TRITON_NO_FILTER + if self._sample_route == _SAMPLE_ROUTE_GUMBEL_TOP_K: + return CUDA_GRAPH_VARIANT_TRITON_TOP_K + if self._sample_route == _SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P: + return CUDA_GRAPH_VARIANT_TRITON_TOP_K_TOP_P + if self._sample_route == _SAMPLE_ROUTE_GUMBEL_TOP_P: + return CUDA_GRAPH_VARIANT_TRITON_TOP_P + return CUDA_GRAPH_VARIANT_DEFAULT + + def sample( + self, + logits_output: LogitsProcessorOutput, + sampling_info: SamplingBatchInfo, + ) -> tuple[torch.Tensor, torch.Tensor]: + logits = nan_guard_logits( + logits_output.next_token_logits, self.config.enable_nan_detection + ) + + if sampling_info.vocab_mask is not None: + sampling_info.apply_vocab_mask( + logits=logits, vocab_mask=sampling_info.vocab_mask + ) + + if sampling_info.is_all_greedy: + batch_next_token_ids = cute_argmax(logits) + else: + offsets_pool = ( + sampling_info.valid_cache_lengths + if sampling_info.valid_cache_lengths is not None + else self._zero_offsets_pool + ) + bs = logits.shape[0] + req_pool_indices = self._req_pool_indices_for_kernels( + sampling_info.req_pool_indices, bs + ) + if self._sample_route == _SAMPLE_ROUTE_GUMBEL_NO_FILTER: + if logits.shape[1] <= _COMPACT_GUMBEL_VOCAB_MAX: + batch_next_token_ids = gumbel_sample_from_pools_compact( + logits, + req_pool_indices, + self._temperature_pool, + self._seed_pool, + offsets_pool, + self._gumbel_out[:bs], + block_size=_COMPACT_GUMBEL_BLOCK_SIZE, + ) + else: + batch_next_token_ids = gumbel_sample_from_pools( + logits, + req_pool_indices, + self._temperature_pool, + self._seed_pool, + offsets_pool, + self._gumbel_local_ids[:bs], + self._gumbel_local_scores[:bs], + self._gumbel_out[:bs], + ) + elif self._sample_route in ( + _SAMPLE_ROUTE_GUMBEL_TOP_K, + _SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P, + ): + batch_next_token_ids = gumbel_sample_top_k_top_p_from_pools( + logits, + req_pool_indices, + self._temperature_pool, + self._top_k_pool, + self._top_p_pool, + self._seed_pool, + offsets_pool, + self._topk_candidate_ids[:bs], + self._topk_candidate_logits[:bs], + self._gumbel_out[:bs], + block_size=_TOP_K_TOP_P_SMALL_BLOCK_SIZE, + top_k_pad=self._top_k_top_p_pad, + ) + elif self._sample_route == _SAMPLE_ROUTE_GUMBEL_TOP_P: + batch_next_token_ids = gumbel_sample_top_p_parallel_from_pools( + logits, + req_pool_indices, + self._temperature_pool, + self._top_p_pool, + self._seed_pool, + offsets_pool, + self._top_p_local_max[:bs], + self._top_p_local_sum[:bs], + self._top_p_local_argmax[:bs], + self._top_p_local_scores[:bs], + self._top_p_local_logits[:bs], + self._top_p_local_ids[:bs], + self._top_p_row_max[:bs], + self._top_p_row_total[:bs], + self._top_p_row_argmax[:bs], + self._top_p_row_candidate_logits[:bs], + self._top_p_row_candidate_ids[:bs], + self._top_p_accepted[:bs], + self._gumbel_out[:bs], + block_size=_TOP_P_PARALLEL_SAMPLE_BLOCK_SIZE, + num_attempts=_TOP_P_PARALLEL_SAMPLE_ATTEMPTS, + ) + else: + batch_next_token_ids = gumbel_sample_from_pools_generic( + logits, + req_pool_indices, + self._temperature_pool, + self._top_k_pool, + self._top_p_pool, + self._seed_pool, + offsets_pool, + self._gumbel_out[:bs], + ) + + sampled = batch_next_token_ids.to(torch.int32) + self.maybe_broadcast(sampled) + + self._write_logprob_outputs(logits_output, logits, sampled) + + return sampled, self._ones_buf[: logits.shape[0]] + + @nvtx_range("sampling:verify", color="yellow") + def verify( + self, + logits_output: LogitsProcessorOutput, + sampling_info: SamplingBatchInfo, + candidates: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + bs = candidates.shape[0] + num_tokens_per_req = candidates.shape[1] + + predict = self._predict_buf[: bs * num_tokens_per_req] + accept_index = ( + self._accept_index_buf[: bs * num_tokens_per_req] + .view(bs, num_tokens_per_req) + .fill_(-1) + ) + accept_length = self._accept_length_buf[:bs] + + logits = nan_guard_logits( + logits_output.next_token_logits, self.config.enable_nan_detection + ) + + if sampling_info.vocab_mask is not None: + sampling_info.apply_vocab_mask( + logits=logits, + vocab_mask=sampling_info.vocab_mask, + ) + + if sampling_info.is_all_greedy: + target_sampled = cute_argmax(logits) + verify_chain_target_sampled( + predicts=predict, + accept_index=accept_index, + accept_token_num=accept_length, + candidates=candidates, + target_sampled=target_sampled, + enable_pdl=pdl_enabled(), + ) + else: + offsets_pool = ( + sampling_info.valid_cache_lengths + if sampling_info.valid_cache_lengths is not None + else self._zero_offsets_pool + ) + req_pool_indices = self._req_pool_indices_for_kernels( + sampling_info.req_pool_indices, bs + ) + if self._sample_route == _SAMPLE_ROUTE_GUMBEL_NO_FILTER: + if logits.shape[1] <= _COMPACT_GUMBEL_VOCAB_MAX: + target_sampled = gumbel_sample_from_pools_compact( + logits, + req_pool_indices, + self._temperature_pool, + self._seed_pool, + offsets_pool, + self._gumbel_verify_out[: bs * num_tokens_per_req], + block_size=_COMPACT_GUMBEL_BLOCK_SIZE, + num_tokens_per_req=num_tokens_per_req, + ) + else: + target_sampled = gumbel_sample_from_pools( + logits, + req_pool_indices, + self._temperature_pool, + self._seed_pool, + offsets_pool, + self._gumbel_verify_local_ids[: bs * num_tokens_per_req], + self._gumbel_verify_local_scores[: bs * num_tokens_per_req], + self._gumbel_verify_out[: bs * num_tokens_per_req], + num_tokens_per_req=num_tokens_per_req, + ) + elif self._sample_route in ( + _SAMPLE_ROUTE_GUMBEL_TOP_K, + _SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P, + ): + rows = bs * num_tokens_per_req + if self._use_qrita_verify_top_k_route(rows, logits.shape[1]): + target_sampled = gumbel_sample_top_k_top_p_qrita_from_pools( + logits, + req_pool_indices, + self._temperature_pool, + self._top_k_pool, + self._top_p_pool, + self._seed_pool, + offsets_pool, + self._qrita_verify_buffer, + self._qrita_percentile_to_std_table, + self._gumbel_verify_out[:rows], + num_tokens_per_req=num_tokens_per_req, + num_programs=min(self._qrita_verify_num_programs, rows), + ) + else: + target_sampled = gumbel_sample_top_k_top_p_from_pools( + logits, + req_pool_indices, + self._temperature_pool, + self._top_k_pool, + self._top_p_pool, + self._seed_pool, + offsets_pool, + self._topk_verify_candidate_ids[:rows], + self._topk_verify_candidate_logits[:rows], + self._gumbel_verify_out[:rows], + block_size=_TOP_K_TOP_P_SMALL_BLOCK_SIZE, + top_k_pad=self._top_k_top_p_pad, + num_tokens_per_req=num_tokens_per_req, + ) + elif self._sample_route == _SAMPLE_ROUTE_GUMBEL_TOP_P: + rows = bs * num_tokens_per_req + target_sampled = gumbel_sample_top_p_parallel_from_pools( + logits, + req_pool_indices, + self._temperature_pool, + self._top_p_pool, + self._seed_pool, + offsets_pool, + self._top_p_local_max[:rows], + self._top_p_local_sum[:rows], + self._top_p_local_argmax[:rows], + self._top_p_local_scores[:rows], + self._top_p_local_logits[:rows], + self._top_p_local_ids[:rows], + self._top_p_row_max[:rows], + self._top_p_row_total[:rows], + self._top_p_row_argmax[:rows], + self._top_p_row_candidate_logits[:rows], + self._top_p_row_candidate_ids[:rows], + self._top_p_accepted[:rows], + self._gumbel_verify_out[:rows], + block_size=_TOP_P_PARALLEL_SAMPLE_BLOCK_SIZE, + num_attempts=_TOP_P_PARALLEL_VERIFY_ATTEMPTS, + num_tokens_per_req=num_tokens_per_req, + ) + else: + target_sampled = gumbel_sample_from_pools_generic( + logits, + req_pool_indices, + self._temperature_pool, + self._top_k_pool, + self._top_p_pool, + self._seed_pool, + offsets_pool, + self._gumbel_verify_out[: bs * num_tokens_per_req], + num_tokens_per_req=num_tokens_per_req, + ) + verify_chain_target_sampled( + predicts=predict, + accept_index=accept_index, + accept_token_num=accept_length, + candidates=candidates, + target_sampled=target_sampled, + enable_pdl=pdl_enabled(), + ) + + accept_length += 1 + + # Rank 0 remains the source of truth for attention-TP agreement. + self.maybe_broadcast(predict, accept_index, accept_length) + + if self.config.enable_output_logprobs: + self._write_logprob_outputs( + logits_output, + logits, + predict, + ) + + return predict, accept_length + + +register_backend("triton", TritonSamplingBackend) diff --git a/python/tokenspeed/runtime/sampling/backends/triton_full.py b/python/tokenspeed/runtime/sampling/backends/triton_full.py new file mode 100644 index 0000000..292d102 --- /dev/null +++ b/python/tokenspeed/runtime/sampling/backends/triton_full.py @@ -0,0 +1,572 @@ +# SPDX-License-Identifier: MIT AND Apache-2.0 +# SPDX-FileCopyrightText: Copyright (c) 2026 LightSeek Foundation +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +# TokenSpeed keeps pool-owned counts and logit-bias state. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from tokenspeed_kernel.ops.sampling.triton import ( + accumulate_counts_inplace, + apply_penalties_logit_bias_inplace, + gumbel_sample_from_pools, + gumbel_sample_from_pools_compact, + gumbel_sample_from_pools_generic, + gumbel_sample_min_p_from_pools, + gumbel_sample_min_p_from_pools_parallel, + gumbel_sample_top_k_top_p_from_pools, + gumbel_sample_top_k_top_p_qrita_from_pools, + gumbel_sample_top_p_parallel_from_pools, + verify_chain_target_sampled, +) + +from tokenspeed.runtime.sampling.backends.base import ( + CUDA_GRAPH_VARIANT_DEFAULT, + SamplingBackend, + SamplingBackendConfig, +) +from tokenspeed.runtime.sampling.backends.triton import ( + _COMPACT_GUMBEL_BLOCK_SIZE, + _COMPACT_GUMBEL_VOCAB_MAX, + _SAMPLE_ROUTE_GUMBEL_NO_FILTER, + _SAMPLE_ROUTE_GUMBEL_TOP_K, + _SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P, + _SAMPLE_ROUTE_GUMBEL_TOP_P, + _TOP_K_TOP_P_PAD, + _TOP_K_TOP_P_SMALL_BLOCK_SIZE, + _TOP_P_PARALLEL_SAMPLE_ATTEMPTS, + _TOP_P_PARALLEL_SAMPLE_BLOCK_SIZE, + _TOP_P_PARALLEL_VERIFY_ATTEMPTS, + TritonSamplingBackend, +) +from tokenspeed.runtime.sampling.registry import register_backend +from tokenspeed.runtime.sampling.utils import nan_guard_logits +from tokenspeed.runtime.utils.nvtx import nvtx_range +from tokenspeed.runtime.utils.pdl import pdl_enabled + +if TYPE_CHECKING: + from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput + from tokenspeed.runtime.sampling.sampling_batch_info import SamplingBatchInfo + from tokenspeed.runtime.sampling.sampling_params import SamplingParams + + +CUDA_GRAPH_VARIANT_TRITON_FULL_MIN_P = "triton_full_min_p" +CUDA_GRAPH_VARIANT_TRITON_FULL_TOP_K_TOP_P_MIN_P = "triton_full_top_k_top_p_min_p" + + +class TritonFullSamplingBackend(TritonSamplingBackend): + """Full sampling backend with TokenSpeed-owned state and Triton kernels.""" + + def __init__(self, config: SamplingBackendConfig) -> None: + super().__init__(config) + if config.max_req_pool_size <= 0 or config.vocab_size <= 0: + raise ValueError( + "TritonFullSamplingBackend requires max_req_pool_size > 0 and " + f"vocab_size > 0; got max_req_pool_size={config.max_req_pool_size}, " + f"vocab_size={config.vocab_size}" + ) + + pool_rows = config.max_req_pool_size + 1 + self._counts = torch.zeros( + (pool_rows, config.vocab_size), + dtype=torch.int32, + device=config.device, + ) + self._logit_bias = torch.zeros( + (pool_rows, config.vocab_size), + dtype=torch.bfloat16, + device=config.device, + ) + self._min_p_pool = torch.zeros( + (pool_rows,), dtype=torch.float32, device=config.device + ) + self._freq_pen_pool = torch.zeros( + (pool_rows,), dtype=torch.bfloat16, device=config.device + ) + self._pres_pen_pool = torch.zeros( + (pool_rows,), dtype=torch.bfloat16, device=config.device + ) + self._rep_pen_pool = torch.full( + (pool_rows,), 1.0, dtype=torch.bfloat16, device=config.device + ) + self._min_p_row_max = torch.empty( + (config.max_bs * config.max_draft_tokens_per_req,), + dtype=torch.float32, + device=config.device, + ) + self._full_has_min_p = True + + def prepare_step( + self, + request_ids: list[str], + request_pool_indices: list[int], + sampling_params_list: list[SamplingParams], + num_tokens_per_req: int = 1, + ) -> None: + super().prepare_step( + request_ids=request_ids, + request_pool_indices=request_pool_indices, + sampling_params_list=sampling_params_list, + num_tokens_per_req=num_tokens_per_req, + ) + self._full_has_min_p = any(float(sp.min_p) > 0.0 for sp in sampling_params_list) + + def prepare_capture(self, bs: int, num_tokens_per_req: int = 1) -> None: + self._full_has_min_p = True + super().prepare_capture(bs=bs, num_tokens_per_req=num_tokens_per_req) + + def prepare_capture_variant( + self, + bs: int, + num_tokens_per_req: int, + variant: str, + ) -> None: + if variant == CUDA_GRAPH_VARIANT_TRITON_FULL_MIN_P: + self._full_has_min_p = True + self._sample_route = _SAMPLE_ROUTE_GUMBEL_NO_FILTER + SamplingBackend.prepare_capture( + self, + bs=bs, + num_tokens_per_req=num_tokens_per_req, + ) + return + if variant == CUDA_GRAPH_VARIANT_TRITON_FULL_TOP_K_TOP_P_MIN_P: + self._full_has_min_p = True + self._sample_route = _SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P + self._top_k_top_p_pad = _TOP_K_TOP_P_PAD + SamplingBackend.prepare_capture( + self, + bs=bs, + num_tokens_per_req=num_tokens_per_req, + ) + return + self._full_has_min_p = variant == CUDA_GRAPH_VARIANT_DEFAULT + super().prepare_capture_variant( + bs=bs, + num_tokens_per_req=num_tokens_per_req, + variant=variant, + ) + + def cuda_graph_capture_variants(self, num_tokens_per_req: int) -> tuple[str, ...]: + return ( + *super().cuda_graph_capture_variants(num_tokens_per_req), + CUDA_GRAPH_VARIANT_TRITON_FULL_MIN_P, + CUDA_GRAPH_VARIANT_TRITON_FULL_TOP_K_TOP_P_MIN_P, + ) + + def cuda_graph_replay_variant(self, num_tokens_per_req: int) -> str: + if self._full_has_min_p: + if self._sample_route == _SAMPLE_ROUTE_GUMBEL_NO_FILTER: + return CUDA_GRAPH_VARIANT_TRITON_FULL_MIN_P + if self._sample_route in ( + _SAMPLE_ROUTE_GUMBEL_TOP_K, + _SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P, + ): + return CUDA_GRAPH_VARIANT_TRITON_FULL_TOP_K_TOP_P_MIN_P + return CUDA_GRAPH_VARIANT_DEFAULT + return super().cuda_graph_replay_variant(num_tokens_per_req) + + def _reset_slot(self, pool_idx: int, sp: SamplingParams) -> None: + super()._reset_slot(pool_idx, sp) + + self._min_p_pool[pool_idx].fill_(float(sp.min_p)) + self._freq_pen_pool[pool_idx].fill_(float(sp.frequency_penalty)) + self._pres_pen_pool[pool_idx].fill_(float(sp.presence_penalty)) + self._rep_pen_pool[pool_idx].fill_(float(sp.repetition_penalty)) + + self._counts[pool_idx].fill_(0) + self._logit_bias[pool_idx].fill_(0.0) + + bias_map = getattr(sp, "logit_bias", None) if sp is not None else None + if bias_map: + vocab = self._logit_bias.shape[1] + raw_ids = [int(tid) for tid in bias_map.keys()] + assert all(0 <= tid < vocab for tid in raw_ids), ( + f"logit_bias contains out-of-vocab token id(s); " + f"vocab_size={vocab}, offending=" + f"{[t for t in raw_ids if not 0 <= t < vocab]}" + ) + token_ids = torch.tensor( + raw_ids, + device=self._logit_bias.device, + dtype=torch.long, + ) + bias_values = torch.tensor( + list(bias_map.values()), + device=self._logit_bias.device, + dtype=torch.bfloat16, + ) + self._logit_bias[pool_idx, token_ids] = bias_values + + def reset_capture_state(self) -> None: + self._counts[0].fill_(0) + + @nvtx_range("sampling:penalties", color="yellow") + def _apply_penalties_and_bias( + self, + logits: torch.Tensor, + req_pool_indices: torch.Tensor, + num_tokens_per_req: int = 1, + ) -> torch.Tensor: + return apply_penalties_logit_bias_inplace( + logits, + req_pool_indices, + self._counts, + self._logit_bias, + self._freq_pen_pool, + self._pres_pen_pool, + self._rep_pen_pool, + num_tokens_per_req=num_tokens_per_req, + ) + + @nvtx_range("sampling:accum_counts", color="yellow") + def _accumulate_counts( + self, + pool_idx: torch.Tensor, + tokens: torch.Tensor, + weights: torch.Tensor, + ) -> None: + accumulate_counts_inplace(self._counts, pool_idx, tokens, weights) + + def _gumbel_sample_full_logits( + self, + logits: torch.Tensor, + req_pool_indices: torch.Tensor, + offsets_pool: torch.Tensor, + out: torch.Tensor, + *, + num_tokens_per_req: int = 1, + ) -> torch.Tensor: + rows = logits.shape[0] + if ( + self._full_has_min_p + and self._sample_route == _SAMPLE_ROUTE_GUMBEL_NO_FILTER + ): + if logits.shape[1] > _COMPACT_GUMBEL_VOCAB_MAX: + local_ids = ( + self._gumbel_local_ids + if num_tokens_per_req == 1 + else self._gumbel_verify_local_ids + ) + local_scores = ( + self._gumbel_local_scores + if num_tokens_per_req == 1 + else self._gumbel_verify_local_scores + ) + return gumbel_sample_min_p_from_pools_parallel( + logits, + req_pool_indices, + self._temperature_pool, + self._min_p_pool, + self._seed_pool, + offsets_pool, + local_ids[:rows], + local_scores[:rows], + self._min_p_row_max[:rows], + out[:rows], + num_tokens_per_req=num_tokens_per_req, + ) + return gumbel_sample_min_p_from_pools( + logits, + req_pool_indices, + self._temperature_pool, + self._min_p_pool, + self._seed_pool, + offsets_pool, + out[:rows], + num_tokens_per_req=num_tokens_per_req, + ) + + if self._sample_route in ( + _SAMPLE_ROUTE_GUMBEL_TOP_K, + _SAMPLE_ROUTE_GUMBEL_TOP_K_TOP_P, + ): + if ( + not self._full_has_min_p + and num_tokens_per_req > 1 + and self._use_qrita_verify_top_k_route(rows, logits.shape[1]) + ): + return gumbel_sample_top_k_top_p_qrita_from_pools( + logits, + req_pool_indices, + self._temperature_pool, + self._top_k_pool, + self._top_p_pool, + self._seed_pool, + offsets_pool, + self._qrita_verify_buffer, + self._qrita_percentile_to_std_table, + out[:rows], + num_tokens_per_req=num_tokens_per_req, + num_programs=min(self._qrita_verify_num_programs, rows), + ) + candidate_ids = ( + self._topk_candidate_ids + if num_tokens_per_req == 1 + else self._topk_verify_candidate_ids + ) + candidate_logits = ( + self._topk_candidate_logits + if num_tokens_per_req == 1 + else self._topk_verify_candidate_logits + ) + return gumbel_sample_top_k_top_p_from_pools( + logits, + req_pool_indices, + self._temperature_pool, + self._top_k_pool, + self._top_p_pool, + self._seed_pool, + offsets_pool, + candidate_ids[:rows], + candidate_logits[:rows], + out[:rows], + min_p_pool=self._min_p_pool if self._full_has_min_p else None, + block_size=_TOP_K_TOP_P_SMALL_BLOCK_SIZE, + top_k_pad=self._top_k_top_p_pad, + num_tokens_per_req=num_tokens_per_req, + ) + + if not self._full_has_min_p: + if self._sample_route == _SAMPLE_ROUTE_GUMBEL_NO_FILTER: + if logits.shape[1] <= _COMPACT_GUMBEL_VOCAB_MAX: + return gumbel_sample_from_pools_compact( + logits, + req_pool_indices, + self._temperature_pool, + self._seed_pool, + offsets_pool, + out[:rows], + block_size=_COMPACT_GUMBEL_BLOCK_SIZE, + num_tokens_per_req=num_tokens_per_req, + ) + local_ids = ( + self._gumbel_local_ids + if num_tokens_per_req == 1 + else self._gumbel_verify_local_ids + ) + local_scores = ( + self._gumbel_local_scores + if num_tokens_per_req == 1 + else self._gumbel_verify_local_scores + ) + return gumbel_sample_from_pools( + logits, + req_pool_indices, + self._temperature_pool, + self._seed_pool, + offsets_pool, + local_ids[:rows], + local_scores[:rows], + out[:rows], + num_tokens_per_req=num_tokens_per_req, + ) + if self._sample_route == _SAMPLE_ROUTE_GUMBEL_TOP_P: + return gumbel_sample_top_p_parallel_from_pools( + logits, + req_pool_indices, + self._temperature_pool, + self._top_p_pool, + self._seed_pool, + offsets_pool, + self._top_p_local_max[:rows], + self._top_p_local_sum[:rows], + self._top_p_local_argmax[:rows], + self._top_p_local_scores[:rows], + self._top_p_local_logits[:rows], + self._top_p_local_ids[:rows], + self._top_p_row_max[:rows], + self._top_p_row_total[:rows], + self._top_p_row_argmax[:rows], + self._top_p_row_candidate_logits[:rows], + self._top_p_row_candidate_ids[:rows], + self._top_p_accepted[:rows], + out[:rows], + block_size=_TOP_P_PARALLEL_SAMPLE_BLOCK_SIZE, + num_attempts=( + _TOP_P_PARALLEL_SAMPLE_ATTEMPTS + if num_tokens_per_req == 1 + else _TOP_P_PARALLEL_VERIFY_ATTEMPTS + ), + num_tokens_per_req=num_tokens_per_req, + ) + + return gumbel_sample_from_pools_generic( + logits, + req_pool_indices, + self._temperature_pool, + self._top_k_pool, + self._top_p_pool, + self._seed_pool, + offsets_pool, + out[:rows], + min_p_pool=self._min_p_pool, + num_tokens_per_req=num_tokens_per_req, + ) + + @nvtx_range("sampling:sample", color="yellow") + def sample( + self, + logits_output: LogitsProcessorOutput, + sampling_info: SamplingBatchInfo, + ) -> tuple[torch.Tensor, torch.Tensor]: + logits = nan_guard_logits( + logits_output.next_token_logits, self.config.enable_nan_detection + ).float() + + if sampling_info.vocab_mask is not None: + sampling_info.apply_vocab_mask( + logits=logits, vocab_mask=sampling_info.vocab_mask + ) + + logits_for_logprobs = ( + logits.clone() if self.config.enable_output_logprobs else None + ) + + req_pool_indices = self._req_pool_indices_for_kernels( + sampling_info.req_pool_indices, logits.shape[0] + ) + logits = self._apply_penalties_and_bias(logits, req_pool_indices) + offsets_pool = ( + sampling_info.valid_cache_lengths + if sampling_info.valid_cache_lengths is not None + else self._zero_offsets_pool + ) + sampled = self._gumbel_sample_full_logits( + logits, + req_pool_indices, + offsets_pool, + self._gumbel_out[: logits.shape[0]], + ).to(torch.int32) + + self.maybe_broadcast(sampled) + + if logits_for_logprobs is not None: + self._write_logprob_outputs( + logits_output, + logits_for_logprobs, + sampled, + ) + + self._accumulate_counts( + req_pool_indices, + sampled, + torch.ones_like(sampled, dtype=torch.int32), + ) + + return sampled, self._ones_buf[: logits.shape[0]] + + @nvtx_range("sampling:verify", color="yellow") + def verify( + self, + logits_output: LogitsProcessorOutput, + sampling_info: SamplingBatchInfo, + candidates: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + bs = candidates.shape[0] + num_tokens_per_req = candidates.shape[1] + + predict = self._predict_buf[: bs * num_tokens_per_req] + accept_index = ( + self._accept_index_buf[: bs * num_tokens_per_req] + .view(bs, num_tokens_per_req) + .fill_(-1) + ) + accept_length = self._accept_length_buf[:bs] + + logits = nan_guard_logits( + logits_output.next_token_logits, self.config.enable_nan_detection + ).float() + + if sampling_info.vocab_mask is not None: + sampling_info.apply_vocab_mask( + logits=logits, + vocab_mask=sampling_info.vocab_mask, + ) + + logits_for_logprobs = ( + logits.clone() if self.config.enable_output_logprobs else None + ) + + req_pool_indices = self._req_pool_indices_for_kernels( + sampling_info.req_pool_indices, bs + ) + logits = self._apply_penalties_and_bias( + logits, + req_pool_indices, + num_tokens_per_req=num_tokens_per_req, + ) + + offsets_pool = ( + sampling_info.valid_cache_lengths + if sampling_info.valid_cache_lengths is not None + else self._zero_offsets_pool + ) + target_sampled = self._gumbel_sample_full_logits( + logits, + req_pool_indices, + offsets_pool, + self._gumbel_verify_out[: bs * num_tokens_per_req], + num_tokens_per_req=num_tokens_per_req, + ) + verify_chain_target_sampled( + predicts=predict, + accept_index=accept_index, + accept_token_num=accept_length, + candidates=candidates.to(torch.int32), + target_sampled=target_sampled, + enable_pdl=pdl_enabled(), + ) + + accept_length += 1 + + self.maybe_broadcast(predict, accept_index, accept_length) + + valid = accept_index >= 0 + safe_positions = accept_index.clamp(min=0).long() + accepted_tokens = predict.long().gather(0, safe_positions.view(-1)) + + pool_idx_expanded = ( + req_pool_indices.unsqueeze(-1).expand(-1, num_tokens_per_req).reshape(-1) + ) + + self._accumulate_counts( + pool_idx_expanded, + accepted_tokens, + valid.reshape(-1).to(torch.int32), + ) + + if logits_for_logprobs is not None: + self._write_logprob_outputs( + logits_output, + logits_for_logprobs, + predict, + ) + + return predict, accept_length + + +register_backend("triton_full", TritonFullSamplingBackend) diff --git a/python/tokenspeed/runtime/sampling/custom_logit_processor.py b/python/tokenspeed/runtime/sampling/custom_logit_processor.py new file mode 100755 index 0000000..1ac8b1d --- /dev/null +++ b/python/tokenspeed/runtime/sampling/custom_logit_processor.py @@ -0,0 +1,58 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import json +from abc import ABC, abstractmethod +from functools import cache +from typing import Any, Self + +import dill +import torch + + +@cache +def _cache_from_str(json_str: str) -> "CustomLogitProcessor": + """Deserialize a json string to a Callable object. + This function is cached to avoid redundant deserialization. + """ + data = json.loads(json_str) + return dill.loads(bytes.fromhex(data["callable"])) + + +class CustomLogitProcessor(ABC): + """Abstract base class for callable functions.""" + + @abstractmethod + def __call__( + self, + logits: torch.Tensor, + custom_param_list: list[dict[str, Any]] | None = None, + ) -> torch.Tensor: + """Define the callable behavior.""" + raise NotImplementedError + + def to_str(self) -> str: + """Serialize the callable function to a JSON-compatible string.""" + return json.dumps({"callable": dill.dumps(self).hex()}) + + @classmethod + def from_str(cls, json_str: str) -> Self: + """Deserialize a callable function from a JSON string.""" + return _cache_from_str(json_str) diff --git a/python/tokenspeed/runtime/sampling/dp_sampling_config.py b/python/tokenspeed/runtime/sampling/dp_sampling_config.py new file mode 100644 index 0000000..94f98e3 --- /dev/null +++ b/python/tokenspeed/runtime/sampling/dp_sampling_config.py @@ -0,0 +1,190 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import dataclasses + +import torch + + +@dataclasses.dataclass(frozen=True) +class DpSamplingSupport: + requested: bool + enabled: bool + infra_supports: bool + drafter_available: bool + backend_supports_verify: bool + tp_size: int + tp_group_set: bool + + +@dataclasses.dataclass(frozen=True) +class DpSamplingTopology: + tp_rank: int + tp_size: int + tp_group: tuple[int, ...] | None + skip_all_gather: bool + tie_word_embeddings: bool = False + + @property + def tp_group_set(self) -> bool: + return self.tp_group is not None + + +@dataclasses.dataclass(frozen=True) +class DpSamplingRuntimeConfig: + enabled: bool = False + vocab_size: int | None = None + max_bucket_bs: int | None = None + min_bs: int | None = None + num_tokens_per_req: int = 1 + topology: DpSamplingTopology | None = None + device: torch.device | str | None = None + + +@dataclasses.dataclass(frozen=True) +class DpSamplingRuntimeLimits: + runtime_vocab_size: int + max_num_seqs: int + data_parallel_size: int + num_tokens_per_req: int + configured_min_bs: int | None + device: torch.device | str + + +def resolve_dp_sampling_support( + *, + requested: bool, + drafter_available: bool, + backend_supports_verify: bool, + topology: DpSamplingTopology, +) -> DpSamplingSupport: + tp_size = int(topology.tp_size) + tp_group_set = topology.tp_group_set + infra_supports = ( + drafter_available and backend_supports_verify and tp_size > 1 and tp_group_set + ) + support = DpSamplingSupport( + requested=bool(requested), + enabled=infra_supports and bool(requested), + infra_supports=infra_supports, + drafter_available=drafter_available, + backend_supports_verify=backend_supports_verify, + tp_size=tp_size, + tp_group_set=tp_group_set, + ) + if support.requested and not support.infra_supports: + raise RuntimeError( + "--dp-sampling was set but Batch-DP spec-verify " + "preconditions are not met: " + f"drafter={support.drafter_available}, " + f"backend_supports_dp_verify={support.backend_supports_verify}, " + f"tp_size={support.tp_size}, " + f"tp_group_set={support.tp_group_set}" + ) + return support + + +def resolve_dp_sampling_runtime( + *, + support: DpSamplingSupport, + lm_head_rows: int, + topology: DpSamplingTopology, + limits: DpSamplingRuntimeLimits, +) -> DpSamplingRuntimeConfig: + if not support.enabled: + return DpSamplingRuntimeConfig( + num_tokens_per_req=limits.num_tokens_per_req, + topology=topology, + device=limits.device, + ) + validate_dp_sampling_lm_head_vocab( + lm_head_rows=lm_head_rows, + vocab_size=limits.runtime_vocab_size, + tp_size=topology.tp_size, + skip_all_gather=topology.skip_all_gather, + tie_word_embeddings=topology.tie_word_embeddings, + ) + dp_vocab_size = int(lm_head_rows) + if not topology.skip_all_gather: + dp_vocab_size *= int(topology.tp_size) + dp_vocab_size = ( + (dp_vocab_size + int(topology.tp_size) - 1) // int(topology.tp_size) + ) * int(topology.tp_size) + max_bs = limits.max_num_seqs // max(limits.data_parallel_size, 1) + max_bucket_bs = ( + (max_bs + topology.tp_size - 1) // topology.tp_size + ) * topology.tp_size + min_bs = ( + int(limits.configured_min_bs) + if limits.configured_min_bs is not None + else 2 * int(topology.tp_size) + ) + if min_bs < 1: + raise ValueError("dp_sampling_min_bs must be >= 1") + return DpSamplingRuntimeConfig( + enabled=True, + vocab_size=dp_vocab_size, + max_bucket_bs=max_bucket_bs, + min_bs=min_bs, + num_tokens_per_req=limits.num_tokens_per_req, + topology=topology, + device=limits.device, + ) + + +def slice_dp_vocab_mask( + vocab_mask: torch.Tensor | None, + *, + full_bs: int, + pad_bs: int, + num_tokens_per_req: int, + shard: slice, +) -> torch.Tensor | None: + if vocab_mask is None: + return None + n = num_tokens_per_req + if pad_bs > full_bs: + vocab_mask = torch.nn.functional.pad( + vocab_mask, + (0, 0, 0, (pad_bs - full_bs) * n), + value=-1, + ) + return vocab_mask.view(pad_bs, n, -1)[shard].reshape(-1, vocab_mask.shape[-1]) + + +def validate_dp_sampling_lm_head_vocab( + *, + lm_head_rows: int, + vocab_size: int, + tp_size: int, + skip_all_gather: bool, + tie_word_embeddings: bool, +) -> None: + if skip_all_gather and int(lm_head_rows) < int(vocab_size): + raise RuntimeError( + "Batch-DP sampling with skip_all_gather requires a replicated/" + "full-vocab LM head. Got a sharded LM head with " + f"lm_head_rows={lm_head_rows}, vocab_size={vocab_size}, " + f"tp_size={tp_size}, tie_word_embeddings={tie_word_embeddings}. " + "Disable --dp-sampling or use a model path that resolves a " + "replicated LM head." + ) diff --git a/python/tokenspeed/runtime/sampling/logits_layout.py b/python/tokenspeed/runtime/sampling/logits_layout.py new file mode 100644 index 0000000..aa2a0ab --- /dev/null +++ b/python/tokenspeed/runtime/sampling/logits_layout.py @@ -0,0 +1,119 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Declarative logits layout plans used by sampling.""" + +from __future__ import annotations + +import dataclasses + +import torch + +from tokenspeed.runtime.distributed.comm_backend import Group +from tokenspeed.runtime.distributed.dp_sampling_comm import DpSamplingComm + + +@dataclasses.dataclass(frozen=True) +class LogitsLayoutPlan: + effective_bs: int + bucket_bs: int + tp_size: int + num_tokens_per_req: int + + +class LogitsLayoutExecutor: + """Executes sampling-provided logits layout plans.""" + + def __init__( + self, + *, + tp_rank: int, + tp_size: int, + tp_group: Group, + max_bucket_bs: int, + num_tokens_per_req: int, + vocab_size: int, + device: torch.device | str, + ) -> None: + self._tp_rank = tp_rank + self._tp_size = tp_size + self._num_tokens_per_req = num_tokens_per_req + self._comm = DpSamplingComm( + tp_size=tp_size, + rank=tp_rank, + group=tp_group, + max_pad_bs=max_bucket_bs, + num_tokens_per_req=num_tokens_per_req, + vocab_size=vocab_size, + logits_dtype=None, + device=device, + ) + + def slice_hidden_states( + self, + hidden_states: torch.Tensor, + plan: LogitsLayoutPlan, + ) -> torch.Tensor: + n = self._tokens_per_req(plan) + rows = hidden_states.shape[0] + if rows % n != 0: + raise ValueError(f"hidden_states have {rows} rows, not divisible by N={n}") + bs = rows // n + if bs != plan.effective_bs: + raise ValueError( + f"hidden_states imply effective_bs={bs}, but logits layout plan has " + f"effective_bs={plan.effective_bs}" + ) + pad_rows = (plan.bucket_bs - plan.effective_bs) * n + if pad_rows > 0: + hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, pad_rows)) + reqs_per_rank = plan.bucket_bs // self._tp_size + start = self._tp_rank * reqs_per_rank * n + return hidden_states[start : start + reqs_per_rank * n] + + def swap_batch_vocab( + self, + local_logits: torch.Tensor, + plan: LogitsLayoutPlan, + ) -> torch.Tensor: + n = self._tokens_per_req(plan) + rows = local_logits.shape[0] + if rows % n != 0: + raise ValueError(f"local logits have {rows} rows, not divisible by N={n}") + bs = rows // n + if bs != plan.effective_bs: + raise ValueError( + f"local logits imply effective_bs={bs}, but logits layout plan has " + f"effective_bs={plan.effective_bs}" + ) + pad_rows = (plan.bucket_bs - plan.effective_bs) * n + if pad_rows > 0: + local_logits = torch.nn.functional.pad(local_logits, (0, 0, 0, pad_rows)) + return self._comm.swap_batch_vocab(local_logits, pad_bs=plan.bucket_bs) + + def _tokens_per_req(self, plan: LogitsLayoutPlan) -> int: + if ( + plan.tp_size != self._tp_size + or plan.num_tokens_per_req != self._num_tokens_per_req + or plan.bucket_bs < plan.effective_bs + or plan.bucket_bs % self._tp_size != 0 + ): + raise RuntimeError("invalid DP logits layout plan") + return plan.num_tokens_per_req diff --git a/python/tokenspeed/runtime/sampling/registry.py b/python/tokenspeed/runtime/sampling/registry.py new file mode 100644 index 0000000..be071b6 --- /dev/null +++ b/python/tokenspeed/runtime/sampling/registry.py @@ -0,0 +1,93 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from tokenspeed_kernel.platform import current_platform + +from tokenspeed.runtime.sampling.backends.base import ( + DEFAULT_RANDOM_SEED, + SamplingBackend, + SamplingBackendConfig, +) + +if TYPE_CHECKING: + from tokenspeed.runtime.utils.server_args import ServerArgs + + +_BACKEND_REGISTRY: dict[str, type[SamplingBackend]] = {} + + +def _get_default_backend_name() -> str: + if current_platform().is_nvidia: + return "flashinfer" + return "greedy" + + +def _resolve_backend_name(server_args: ServerArgs) -> str: + return server_args.sampling_backend or _get_default_backend_name() + + +def register_backend(name: str, cls: type[SamplingBackend]) -> None: + _BACKEND_REGISTRY[name] = cls + + +def create_sampling_backend( + server_args: ServerArgs, + *, + max_bs: int, + max_draft_tokens_per_req: int, + device: str, + random_seed: int = DEFAULT_RANDOM_SEED, + max_req_pool_size: int = 0, + vocab_size: int = 0, + tp_group: tuple[int, ...] | None = None, +) -> SamplingBackend: + # Trigger concrete-backend registration on first use. + from tokenspeed.runtime.sampling.backends import flashinfer as _fi # noqa: F401 + from tokenspeed.runtime.sampling.backends import ( # noqa: F401 + flashinfer_full as _ff, + ) + from tokenspeed.runtime.sampling.backends import greedy as _g # noqa: F401 + from tokenspeed.runtime.sampling.backends import triton as _t # noqa: F401 + from tokenspeed.runtime.sampling.backends import triton_full as _tf # noqa: F401 + + name = _resolve_backend_name(server_args) + if name not in _BACKEND_REGISTRY: + raise ValueError( + f"Unknown sampling backend: {name!r}. " + f"Available: {list(_BACKEND_REGISTRY)}" + ) + cls = _BACKEND_REGISTRY[name] + + return cls( + SamplingBackendConfig.from_server_args( + server_args, + max_bs=max_bs, + max_draft_tokens_per_req=max_draft_tokens_per_req, + device=device, + random_seed=random_seed, + max_req_pool_size=max_req_pool_size, + vocab_size=vocab_size, + tp_group=tp_group, + ) + ) diff --git a/python/tokenspeed/runtime/sampling/sampling_batch_info.py b/python/tokenspeed/runtime/sampling/sampling_batch_info.py new file mode 100755 index 0000000..b3abc56 --- /dev/null +++ b/python/tokenspeed/runtime/sampling/sampling_batch_info.py @@ -0,0 +1,140 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import dataclasses +import threading +from collections.abc import Callable +from typing import TYPE_CHECKING + +import torch + +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +if TYPE_CHECKING: + from tokenspeed.runtime.engine.schedule_batch import ScheduleBatch + + +@dataclasses.dataclass +class SamplingBatchInfo: + # Basic batched sampling params. Disaggregated decode populates these via + # from_schedule_batch. The standard hot path leaves them None; sampling + # backends gather params from their own pool-indexed buffers. + temperatures: torch.Tensor | None = None + top_ps: torch.Tensor | None = None + top_ks: torch.Tensor | None = None + min_ps: torch.Tensor | None = None + + # Whether all requests use greedy sampling + is_all_greedy: bool = False + + # Masking tensors for grammar-guided structured outputs + vocab_size: int = 0 + grammars: list | None = None + vocab_mask: torch.Tensor | None = None + # Backend-specific in-place fn ``(logits, vocab_mask) -> None``, + # bound by ``capturable_grammar.bind_grammar_mask_buf`` so the + # captured sampler can apply the bitmask without branching on + # backend. + apply_vocab_mask: Callable[[torch.Tensor, torch.Tensor], None] | None = None + + # An event used for overlap schedule + sampling_info_done: threading.Event | None = None + + # int64[bs] — req_pool_idx per batch row. Sampling backends gather + # their pool-indexed scalar buffers (temperature / top_k / top_p / + # seeds / penalties / logit_bias / counts) against this index. + req_pool_indices: torch.Tensor | None = None + + # int32[pool_rows] — RuntimeStates.valid_cache_lengths, read-only + # reference. Sampling backends derive the per-request Philox offset + # from `valid_cache_lengths.index_select(0, req_pool_indices)`; + # carrying the reference rather than the gathered view keeps the + # index_select inside the captured graph. + valid_cache_lengths: torch.Tensor | None = None + + # Device + device: str = "cuda" + + def __getitem__(self, s: slice) -> SamplingBatchInfo: + """Row-slice batch-indexed fields; pool/scalar fields pass through. + + Used by hybrid-batch samplers (MIXED + spec-dec) that apply + different sampler ops to a prefix vs suffix of rows. Only ``slice`` + is supported — int indexing would yield 0-dim tensors and break + downstream gathers. + + ``is_all_greedy`` is inherited from the parent; when ``top_ks`` is + populated the slice refines it from the sliced tensor (one GPU + sync, only on the disagg slice path). + """ + if not isinstance(s, slice): + raise TypeError( + f"SamplingBatchInfo only supports slice indexing, got {type(s).__name__}" + ) + + def _slice(t): + return t[s] if t is not None else None + + return dataclasses.replace( + self, + temperatures=_slice(self.temperatures), + top_ps=_slice(self.top_ps), + top_ks=_slice(self.top_ks), + min_ps=_slice(self.min_ps), + is_all_greedy=self.is_all_greedy, + req_pool_indices=_slice(self.req_pool_indices), + vocab_mask=_slice(self.vocab_mask), + grammars=_slice(self.grammars), + ) + + @classmethod + def from_schedule_batch( + cls, batch: ScheduleBatch, vocab_size: int + ) -> SamplingBatchInfo: + reqs = batch.reqs + device = batch.device + temperatures = torch.tensor( + [r.sampling_params.temperature for r in reqs], dtype=torch.float + ).to(device, non_blocking=True) + top_ps = torch.tensor( + [r.sampling_params.top_p for r in reqs], dtype=torch.float + ).to(device, non_blocking=True) + top_ks = torch.tensor( + [r.sampling_params.top_k for r in reqs], dtype=torch.int32 + ).to(device, non_blocking=True) + min_ps = torch.tensor( + [r.sampling_params.min_p for r in reqs], dtype=torch.float + ).to(device, non_blocking=True) + + ret = cls( + temperatures=temperatures, + top_ps=top_ps, + top_ks=top_ks, + min_ps=min_ps, + is_all_greedy=all(r.sampling_params.top_k <= 1 for r in reqs), + vocab_size=vocab_size, + device=device, + ) + return ret diff --git a/python/tokenspeed/runtime/sampling/sampling_params.py b/python/tokenspeed/runtime/sampling/sampling_params.py new file mode 100755 index 0000000..f8fcfc0 --- /dev/null +++ b/python/tokenspeed/runtime/sampling/sampling_params.py @@ -0,0 +1,243 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Sampling parameters for text generation.""" + +import zlib +from typing import Any + +_SAMPLING_EPS = 1e-6 + +# Sentinel for "top_k is disabled" (sample from whole vocab). We rewrite +# top_k=-1 (API convention) to this value so downstream code can pass it +# unchanged to top_k kernels that expect a positive cutoff. +_TOP_K_DISABLED = 1 << 30 + +# Upper bound the fused top-k + top-p kernel sorts in its on-chip top-K +# branch. Requests with a finite top_k above this would silently fall through +# to the top-p-only branch, so reject them at request time. Must stay in sync +# with K_TOPK_MAX in fused_topk_topp.h. +_TOP_K_FUSED_MAX = 128 + + +class SamplingParams: + """ + The sampling parameters. + + See docs/backend/sampling_params.md or + https://docs.tokenspeed.ai/backend/sampling_params.html + for the documentation. + """ + + def __init__( + self, + max_new_tokens: int | None = None, + stop: str | list[str] | None = None, + stop_token_ids: list[int] | None = 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, + repetition_penalty: float = 1.0, + min_new_tokens: int = 0, + json_schema: str | None = None, + regex: str | None = None, + ebnf: str | None = None, + structural_tag: str | None = None, + ignore_eos: bool = False, + skip_special_tokens: bool = True, + spaces_between_special_tokens: bool = True, + no_stop_trim: bool = False, + thinking_budget: int | None = None, + custom_params: dict[str, Any] | None = None, + stream_interval: int | None = None, + logit_bias: dict[str, float] | None = None, + seed: int | None = None, + # vLLM-style output logprobs. None = off; 0 = the sampled (generated) + # token's logprob at each output position. Other values are rejected by + # verify(). + logprobs: int | None = None, + # OpenAI-compat: `n` is a request-level fanout (number of choices) + # that the serving layer forwards on every sampling_params dict. + # TokenSpeed does not multiplex a single request into n completions, + # so accept and ignore. + n: int = 1, + ) -> None: + self.max_new_tokens = max_new_tokens + self.stop_strs = stop + if stop_token_ids: + self.stop_token_ids = set(stop_token_ids) + else: + self.stop_token_ids = None + self.temperature = temperature + self.top_p = top_p + self.top_k = top_k + self.min_p = min_p + self.frequency_penalty = frequency_penalty + self.presence_penalty = presence_penalty + self.repetition_penalty = repetition_penalty + self.min_new_tokens = min_new_tokens + self.regex = regex + self.json_schema = json_schema + self.ebnf = ebnf + self.structural_tag = structural_tag + self.ignore_eos = ignore_eos + self.skip_special_tokens = skip_special_tokens + self.spaces_between_special_tokens = spaces_between_special_tokens + self.no_stop_trim = no_stop_trim + self.custom_params = custom_params + self.thinking_budget = thinking_budget + self.stream_interval = stream_interval + self.logit_bias = logit_bias + self.seed = seed + self.logprobs = logprobs + + # Process some special cases + if self.temperature < _SAMPLING_EPS: + # top_k = 1 means greedy sampling + self.temperature = 1.0 + self.top_k = 1 + if self.top_k == -1: + self.top_k = _TOP_K_DISABLED + + def verify(self, vocab_size: int) -> None: + if self.temperature < 0.0: + raise ValueError( + f"temperature must be non-negative, got {self.temperature}." + ) + if not 0.0 < self.top_p <= 1.0: + raise ValueError(f"top_p must be in (0, 1], got {self.top_p}.") + if not 0.0 <= self.min_p <= 1.0: + raise ValueError(f"min_p must be in [0, 1], got {self.min_p}.") + if self.top_k < -1 or self.top_k == 0: + raise ValueError( + f"top_k must be -1 (disable), or at least 1, " f"got {self.top_k}." + ) + if self.top_k != _TOP_K_DISABLED and self.top_k >= _TOP_K_FUSED_MAX: + raise ValueError( + f"top_k must be < {_TOP_K_FUSED_MAX} (fused kernel limit) " + f"or -1 (disable), got {self.top_k}." + ) + if not -2.0 <= self.frequency_penalty <= 2.0: + raise ValueError( + "frequency_penalty must be in [-2, 2], got " + f"{self.frequency_penalty}." + ) + if not -2.0 <= self.presence_penalty <= 2.0: + raise ValueError( + "presence_penalty must be in [-2, 2], got " f"{self.presence_penalty}." + ) + if not 0.0 <= self.repetition_penalty <= 2.0: + raise ValueError( + "repetition_penalty must be in (0, 2], got " + f"{self.repetition_penalty}." + ) + if not 0 <= self.min_new_tokens: + raise ValueError( + f"min_new_tokens must be in (0, max_new_tokens], got " + f"{self.min_new_tokens}." + ) + if self.max_new_tokens is not None: + if self.max_new_tokens < 0: + raise ValueError( + f"max_new_tokens must be at least 0, got {self.max_new_tokens}." + ) + if not self.min_new_tokens <= self.max_new_tokens: + raise ValueError( + f"min_new_tokens must be in (0, max_new_tokens({self.max_new_tokens})], got " + f"{self.min_new_tokens}." + ) + if self.logit_bias is not None: + for token_id in self.logit_bias: + if not 0 <= int(token_id) < vocab_size: + raise ValueError( + f"logit_bias must has keys in [0, {vocab_size - 1}], got " + f"{token_id}." + ) + + if self.logprobs is not None and self.logprobs != 0: + # Only the sampled token's logprob (logprobs=0) is materialized; + # top-k (>0) and full-vocab (-1) output logprobs are not supported. + raise ValueError( + f"logprobs={self.logprobs} is not supported; use logprobs=0 " + "(the sampled token's logprob)." + ) + + grammars = [ + self.json_schema, + self.regex, + self.ebnf, + ] # since mutually exclusive, only one can be set + if sum(x is not None for x in grammars) > 1: + raise ValueError("Only one of regex, json_schema, or ebnf can be set.") + + def requested_features(self) -> "set[str]": + """Return the set of backend-facing feature names this request needs. + + `temperature`, `top_k`, `top_p`, `min_p` each appear only when the + corresponding field is not at its neutral default. Used by + SamplingBackend.register() to reject requests asking for features + the active backend does not implement.""" + out: set[str] = set() + if abs(self.temperature - 1.0) > _SAMPLING_EPS: + out.add("temperature") + # top_k=_TOP_K_DISABLED and top_k=1 (greedy short-circuit from __init__) are neutral. + if self.top_k != _TOP_K_DISABLED and self.top_k != 1: + out.add("top_k") + if self.top_p < 1.0: + out.add("top_p") + if self.min_p > 0.0: + out.add("min_p") + if self.frequency_penalty != 0.0: + out.add("frequency_penalty") + if self.presence_penalty != 0.0: + out.add("presence_penalty") + if self.repetition_penalty != 1.0: + out.add("repetition_penalty") + if self.logit_bias: + out.add("logit_bias") + return out + + def resolve_seed(self, rid: str) -> None: + """If the caller didn't supply a seed, derive one deterministically + from rid. Called at the single request-materialization point so all + TP/DP ranks agree on the seed.""" + if self.seed is None: + self.seed = zlib.crc32(rid.encode("utf-8")) & 0xFFFFFFFF + + def normalize(self, tokenizer) -> None: + # Process stop strings + if self.stop_strs is None: + self.stop_strs = [] + self.stop_str_max_len = 0 + else: + if isinstance(self.stop_strs, str): + self.stop_strs = [self.stop_strs] + + stop_str_max_len = 0 + for stop_str in self.stop_strs: + if tokenizer is not None: + stop_str_ids = tokenizer.encode(stop_str, add_special_tokens=False) + stop_str_max_len = max(stop_str_max_len, len(stop_str_ids)) + else: + stop_str_max_len = max(stop_str_max_len, len(stop_str)) + self.stop_str_max_len = stop_str_max_len diff --git a/python/tokenspeed/runtime/sampling/utils.py b/python/tokenspeed/runtime/sampling/utils.py new file mode 100644 index 0000000..8c71ea5 --- /dev/null +++ b/python/tokenspeed/runtime/sampling/utils.py @@ -0,0 +1,86 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import torch +from tokenspeed_kernel.torch_compile import get_compiler_backend + +from tokenspeed.runtime.utils import crash_on_warnings, get_colorful_logger + +logger = get_colorful_logger(__name__) + +# Smallest positive value per dtype, used as the lower bound for `uniform_` +# draws that feed rejection-sampling kernels. A coin of exact 0 silently +# accepts a zero-probability draft in `chain_speculative_sampling_target_only` +# (the kernel condition `coin <= target_prob / threshold_acc` reduces to +# `0 <= 0`), so the coin must be strictly positive. +COIN_EPS = { + torch.float32: torch.finfo(torch.float32).tiny, + torch.bfloat16: torch.finfo(torch.bfloat16).tiny, +} + + +def coin_eps(dtype: torch.dtype) -> float: + """Lower bound for uniform coin draws of the given dtype. See COIN_EPS.""" + return COIN_EPS[dtype] + + +def nan_guard_logits( + logits: torch.Tensor, + enable_nan_detection: bool, +) -> torch.Tensor: + """Replace NaNs with -1e5 and optionally crash; no-op when detection is disabled.""" + if not enable_nan_detection: + return logits + + if not torch.any(torch.isnan(logits)): + return 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.") + return logits + + +@torch.compile(dynamic=True, backend=get_compiler_backend()) +def gather_token_logprobs_torch( + logits: torch.Tensor, + tokens: torch.Tensor, +) -> torch.Tensor: + """Per-row log_softmax(logits)[tokens]. Fuses cast + online softmax + gather + into one Triton kernel sequence so the full [B, V] log_softmax matrix is + never materialized.""" + raw_logprobs = torch.log_softmax(logits.float(), dim=-1) + return raw_logprobs.gather(-1, tokens.unsqueeze(-1)).squeeze(-1) + + +@torch.compile(dynamic=True, backend=get_compiler_backend()) +def top_p_normalize_probs_torch( + probs: torch.Tensor, + top_ps: torch.Tensor, +) -> torch.Tensor: + """Pure-torch nucleus renorm — used by the prefill-logprob path.""" + 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) diff --git a/python/tokenspeed/runtime/spec_decode/algorithm.py b/python/tokenspeed/runtime/spec_decode/algorithm.py new file mode 100755 index 0000000..259012b --- /dev/null +++ b/python/tokenspeed/runtime/spec_decode/algorithm.py @@ -0,0 +1,50 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +from enum import IntEnum, auto + + +class SpeculativeAlgorithm(IntEnum): + NONE = auto() + EAGLE3 = auto() + MTP = auto() + DFLASH = auto() + + def is_none(self) -> bool: + return self == SpeculativeAlgorithm.NONE + + def needs_draft_decode_prealloc(self) -> bool: + return self in (SpeculativeAlgorithm.EAGLE3, SpeculativeAlgorithm.MTP) + + @staticmethod + def from_string(name: str | None) -> SpeculativeAlgorithm: + if name is None: + return SpeculativeAlgorithm.NONE + name_map = { + "EAGLE3": SpeculativeAlgorithm.EAGLE3, + "MTP": SpeculativeAlgorithm.MTP, + "DFLASH": SpeculativeAlgorithm.DFLASH, + } + try: + return name_map[name.upper()] + except KeyError as exc: + raise ValueError(f"Unknown speculative algorithm: {name}") from exc diff --git a/python/tokenspeed/runtime/spec_decode/eagle.py b/python/tokenspeed/runtime/spec_decode/eagle.py new file mode 100755 index 0000000..a0a0a53 --- /dev/null +++ b/python/tokenspeed/runtime/spec_decode/eagle.py @@ -0,0 +1,253 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import dataclasses + +import torch +import triton +import triton.language as tl + +from tokenspeed.runtime.execution.forward_batch_info import CaptureHiddenMode +from tokenspeed.runtime.layers.attention.utils import ( + create_flashinfer_kv_indices_triton, +) +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + + +@dataclasses.dataclass +class EagleDraftInput: + # The inputs for decode + # shape: (b, topk) + topk_p: torch.Tensor | None = None + topk_index: torch.Tensor | None = None + # shape: (b, hidden_size) + hidden_states: torch.Tensor | None = None + capture_hidden_mode: CaptureHiddenMode = CaptureHiddenMode.FULL + + # Inputs for extend + # shape: (b,) + verified_id: torch.Tensor | None = None + accept_length: torch.Tensor | None = None + accept_length_cpu: list[int] | None = None + accept_index: torch.Tensor | None = None + + # Inputs for the attention backends + # shape: (b + 1,) + kv_indptr: torch.Tensor | None = None + kv_indices: torch.Tensor | None = None + + # For draft extend fast plan + qo_indptr_cpu: torch.Tensor | None = None + kv_indptr_cpu: torch.Tensor | None = None + kv_indices_for_extend: torch.Tensor | None = None + kv_len_arr_cpu: torch.Tensor | None = None + + draft_token_num: int = 0 + + def set_input_ids( + self, + input_ids: torch.Tensor, + draft_input_ids: torch.Tensor, + extend_seq_lens: torch.Tensor, + ) -> None: + pt = 0 + for i, extend_seq_len in enumerate(extend_seq_lens): + cur_input_ids = draft_input_ids[i] + if cur_input_ids[-1] == -1: + cur_input_ids[-1] = self.verified_id[i] + input_ids[pt : pt + extend_seq_len] = cur_input_ids + + pt += extend_seq_len + + def prepare_extend_after_decode(self, batch_size: int) -> torch.Tensor: + new_verified_id = torch.empty_like(self.accept_length, dtype=torch.long) + create_extend_spec_info[(batch_size,)]( + self.verified_id, + new_verified_id, + self.accept_length, + self.draft_token_num, + ) + # Extract the last accepted token for each request + self.verified_id = new_verified_id + return self.verified_id + + def filter_batch(self, new_indices: torch.Tensor) -> None: + if self.topk_p is not None: + self.topk_p = self.topk_p[: len(new_indices)] + self.topk_index = self.topk_index[: len(new_indices)] + self.hidden_states = self.hidden_states[: len(new_indices)] + self.verified_id = self.verified_id[: len(new_indices)] + + def merge_batch(self, spec_info: EagleDraftInput) -> None: + if self.hidden_states is None: + self.hidden_states = spec_info.hidden_states + self.verified_id = spec_info.verified_id + self.topk_p = spec_info.topk_p + self.topk_index = spec_info.topk_index + return + if spec_info.hidden_states is None: + return + self.hidden_states = torch.cat( + [self.hidden_states, spec_info.hidden_states], dim=0 + ) + self.verified_id = torch.cat([self.verified_id, spec_info.verified_id], dim=0) + if self.topk_p is not None and spec_info.topk_p is not None: + self.topk_p = torch.cat([self.topk_p, spec_info.topk_p]) + self.topk_index = torch.cat([self.topk_index, spec_info.topk_index]) + + +@dataclasses.dataclass +class EagleDraftOutput: + """ + Both prefill and decode batches end with draft. Used to store the previous draft's information, + to construct verify's input at the next decode + + Args: + last_verified_ids: + """ + + last_verified_ids: torch.Tensor + token_list: torch.Tensor + + def filter_batch(self, keep_indices: torch.Tensor) -> None: + # 1. chunked prefill + # 2. retract + # 3. Check finished when updating running and getting new + self.last_verified_ids = self.last_verified_ids[keep_indices] + self.token_list = self.token_list[keep_indices, :] + + def merge_batch(self, spec_info: EagleDraftOutput) -> None: + if spec_info.last_verified_ids is None: + return + if self.last_verified_ids is None: + # May reach here when all requests in running batch are finished + self.last_verified_ids = spec_info.last_verified_ids + self.token_list = spec_info.token_list + return + self.last_verified_ids = torch.cat( + [self.last_verified_ids, spec_info.last_verified_ids] + ) + self.token_list = torch.cat([self.token_list, spec_info.token_list], dim=0) + + +@triton.jit +def create_extend_spec_info( + verified_id, # padded verified id + new_verified_id, + accept_length_ptr, + spec_num_tokens: int, +): + pid = tl.program_id(axis=0) + accept_len = tl.load(accept_length_ptr + pid) + last_verified_id = tl.load(verified_id + pid * spec_num_tokens + accept_len) + tl.store(accept_length_ptr + pid, accept_len + 1) + tl.store(new_verified_id + pid, last_verified_id) + + +@triton.jit +def assign_req_to_token_pool( + req_pool_indices, + req_to_token, + start_offset, + end_offset, + out_cache_loc, + pool_len: tl.constexpr, + bs_upper: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 32 + pid = tl.program_id(axis=0) + kv_start = tl.load(start_offset + pid) + kv_end = tl.load(end_offset + pid) + token_pool = req_to_token + tl.load(req_pool_indices + pid) * pool_len + + length_offset = tl.arange(0, bs_upper) + start = tl.load(start_offset + length_offset, mask=length_offset < pid) + end = tl.load(end_offset + length_offset, mask=length_offset < pid) + out_offset = tl.sum(end - start, axis=0) + + out_cache_ptr = out_cache_loc + out_offset + + save_offset = tl.arange(0, BLOCK_SIZE) + kv_start + load_offset = tl.arange(0, BLOCK_SIZE) + + num_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE) + for _ in range(num_loop): + mask = save_offset < kv_end + data = tl.load(out_cache_ptr + load_offset, mask=mask) + tl.store(token_pool + save_offset, data, mask=mask) + save_offset += BLOCK_SIZE + load_offset += BLOCK_SIZE + + +def generate_attn_arg_prefill( + draft_token_num: int, + req_pool_indices: torch.Tensor, + paged_kernel_lens: torch.Tensor, + req_to_token: torch.Tensor, + kv_indices_buf: torch.Tensor | None = None, + draft_decode_step: int | None = None, +): + batch_size = req_pool_indices.shape[0] + if draft_decode_step is not None: + qo_indptr = torch.arange( + 0, + (1 + batch_size), + step=1, + dtype=torch.int32, + device="cuda", + ) + else: + qo_indptr = torch.arange( + 0, + (1 + batch_size) * draft_token_num, + step=draft_token_num, + dtype=torch.int32, + device="cuda", + ) + + cum_kv_seq_len = torch.zeros((batch_size + 1,), dtype=torch.int32, device="cuda") + + if draft_decode_step is None: + paged_kernel_lens = paged_kernel_lens + draft_token_num + else: + paged_kernel_lens = paged_kernel_lens + draft_decode_step + 1 + + torch.cumsum(paged_kernel_lens, dim=0, out=cum_kv_seq_len[1:]) + if kv_indices_buf is not None: + kv_indices = kv_indices_buf + else: + # Prevent kv_indices out of bounds in large steps + kv_indices = torch.empty( + cum_kv_seq_len[-1] + 256, dtype=torch.int32, device="cuda" + ) + create_flashinfer_kv_indices_triton[(batch_size,)]( + req_to_token, + req_pool_indices, + paged_kernel_lens, + cum_kv_seq_len, + None, + kv_indices, + req_to_token.size(1), + ) + return kv_indices, cum_kv_seq_len, qo_indptr, None diff --git a/python/tokenspeed/runtime/utils/__init__.py b/python/tokenspeed/runtime/utils/__init__.py new file mode 100755 index 0000000..251cb36 --- /dev/null +++ b/python/tokenspeed/runtime/utils/__init__.py @@ -0,0 +1,21 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from tokenspeed.runtime.utils.common import * # noqa: F403 diff --git a/python/tokenspeed/runtime/utils/common.py b/python/tokenspeed/runtime/utils/common.py new file mode 100755 index 0000000..356b6eb --- /dev/null +++ b/python/tokenspeed/runtime/utils/common.py @@ -0,0 +1,1215 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Common utilities.""" + +from __future__ import annotations + +import asyncio +import dataclasses +import functools +import io +import ipaddress +import json +import logging +import os +import pickle +import random +import re +import resource +import shutil +import subprocess +import tempfile +import uuid +from collections import OrderedDict +from collections.abc import Callable, Sequence +from contextlib import contextmanager +from dataclasses import dataclass +from functools import lru_cache +from io import BytesIO +from multiprocessing.reduction import ForkingPickler +from pathlib import Path +from typing import ( + Any, + Generic, + Literal, + Protocol, + TypeVar, +) +from urllib.parse import unquote, urlparse + +import numpy as np +import psutil +import pybase64 +import requests +import torch +import torch.distributed +import torch.distributed as dist +import triton +import zmq +from fastapi.responses import ORJSONResponse +from PIL import Image +from pydantic import BaseModel +from starlette.routing import Mount +from tokenspeed_kernel.platform import current_platform + +from tokenspeed.runtime.metrics.func_timer import enable_func_timer + +logger = logging.getLogger(__name__) + +time_infos = {} + + +_warned_bool_env_var_keys = set() + + +def get_bool_env_var(name: str, default: str = "false") -> bool: + # Runtime env helpers still read a few legacy keys directly until the + # central env module owns all boolean parsing. + value = os.getenv(name, default) + value = value.lower() + + truthy_values = ("true", "1") + falsy_values = ("false", "0") + + if (value not in truthy_values) and (value not in falsy_values): + if value not in _warned_bool_env_var_keys: + logger.warning( + "get_bool_env_var(%s) see non-understandable value=%s and treat as false", + name, + value, + ) + _warned_bool_env_var_keys.add(value) + + return value in truthy_values + + +@lru_cache(maxsize=1) +def get_device_module(): + """Get the device module (cuda, hip, etc.) based on the current device.""" + return torch.get_device_module() + + +def maybe_inference_mode(): + from tokenspeed.runtime.utils.env import envs + + if envs.TOKENSPEED_ENABLE_TORCH_INFERENCE_MODE.get(): + return torch.inference_mode() + else: + return torch.no_grad() + + +def maybe_set_numa_aware_cpu_affinity(device_id: int) -> None: + """Pin the current process to ``device_id``'s NUMA-local CPU set. + + NVIDIA-only optimization. No-op if the env var is False, the platform is not + NVIDIA, or the process already has a constrained affinity (e.g., taskset). + """ + from tokenspeed.runtime.utils.env import envs + + if not envs.TOKENSPEED_NUMA_AWARE_WORKER_AFFINITY.get(): + return + platform = current_platform() + if not platform.is_nvidia: + return + + proc = psutil.Process() + if proc.cpu_affinity() != list(range(psutil.cpu_count())): + return + + if device_id >= len(platform.numa_cpu_affinity): + return + + cpu_affinity = platform.numa_cpu_affinity[device_id] + if not cpu_affinity: + return + + proc.cpu_affinity(list(cpu_affinity)) + logger.info( + "Worker process %s pinned to %s NUMA-local CPUs for device %s.", + proc.pid, + len(cpu_affinity), + device_id, + ) + + +def get_available_gpu_memory( + device, gpu_id, distributed=False, empty_cache=True, cpu_group=None +): + """ + Get available memory for cuda:gpu_id device. + When distributed is True, the available memory is the minimum available memory of all GPUs. + """ + if device == "cuda": + num_gpus = torch.cuda.device_count() + if gpu_id >= num_gpus: + raise ValueError(f"gpu_id={gpu_id} must be less than num_gpus={num_gpus}.") + + if torch.cuda.current_device() != gpu_id: + logger.debug( + "Current device is not %s, but %s, which may cause useless " + "memory allocation for torch CUDA context.", + gpu_id, + torch.cuda.current_device(), + ) + + if empty_cache: + torch.cuda.empty_cache() + free_gpu_memory, _ = torch.cuda.mem_get_info(gpu_id) + + if distributed: + tensor = torch.tensor(free_gpu_memory, dtype=torch.float32) + torch.distributed.all_reduce( + tensor, op=torch.distributed.ReduceOp.MIN, group=cpu_group + ) + free_gpu_memory = tensor.item() + + return free_gpu_memory / (1 << 30) + + +def is_pin_memory_available() -> bool: + return torch.cuda.is_available() + + +class LayerFn(Protocol): + def __call__(self, idx: int, prefix: str) -> torch.nn.Module: ... + + +def make_layers( + num_hidden_layers: int, + layer_fn: LayerFn, + prefix: str = "", +) -> torch.nn.ModuleList: + """Make a list of layers with the given layer function""" + start_layer = 0 + end_layer = num_hidden_layers + modules = torch.nn.ModuleList( + [ + layer_fn(idx=idx, prefix=add_prefix(idx, prefix)) + for idx in range(start_layer, end_layer) + ] + ) + return modules + + +def set_random_seed(seed: int) -> None: + """Set the random seed for all libraries.""" + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + + +@dataclass +class ImageData: + url: str + detail: Literal["auto", "low", "high"] | None = "auto" + + +image_extension_names = (".png", ".jpg", ".jpeg", ".webp", ".gif") + + +def is_jpeg_with_cuda(image_bytes: bytes = b"", gpu_image_decode: bool = True) -> bool: + """ + Check three conditions: + 1. whether CUDA is available. + 2. whether input is recognized as JPEG. + 3. whether GPU image decode is enabled. + """ + if not current_platform().is_nvidia or not gpu_image_decode: + return False + if image_bytes != b"": + return image_bytes.startswith(b"\xff\xd8") and image_bytes.endswith(b"\xff\xd9") + return False + + +def _load_image( + image_bytes: bytes = b"", + image_file: str = "", + gpu_image_decode: bool = True, +) -> torch.Tensor | Image.Image: + """ + Try to decode JPEG with nvJPEG on GPU and return a torch device tensor, + otherwise fallback to decode with PIL on CPU and return a PIL Image. + """ + if image_file != "": + image_bytes = get_image_bytes(image_file) + if is_jpeg_with_cuda(image_bytes, gpu_image_decode): + try: + from torchvision.io import decode_jpeg + + encoded_image = torch.frombuffer(image_bytes, dtype=torch.uint8) + image_tensor = decode_jpeg(encoded_image, device="cuda") + return image_tensor + except Exception as exc: + logger.warning( + f"Failed to decode JPEG on GPU, falling back to CPU. Error: {exc}" + ) + return Image.open(BytesIO(image_bytes)) + + +def load_image( + image_file: Image.Image | str | ImageData | bytes, + gpu_image_decode: bool = True, +) -> tuple[torch.Tensor | Image.Image, tuple[int, int] | None]: + """ + Load image from multiple input formats, including: + ImageData, PIL Image, bytes, URL, file path, file:// URL, data URL, or base64. + """ + if isinstance(image_file, ImageData): + image_file = image_file.url + + image = None + image_size: tuple[int, int] | None = None + if isinstance(image_file, Image.Image): + image = image_file + image_size = (image.width, image.height) + elif isinstance(image_file, bytes): + image = _load_image(image_bytes=image_file, gpu_image_decode=gpu_image_decode) + elif isinstance(image_file, str) and image_file.startswith(("http://", "https://")): + image = _load_image(image_file=image_file, gpu_image_decode=gpu_image_decode) + elif isinstance(image_file, str) and image_file.startswith("file://"): + image = _load_image( + image_file=unquote(urlparse(image_file).path), + gpu_image_decode=gpu_image_decode, + ) + elif isinstance(image_file, str) and image_file.lower().endswith( + image_extension_names + ): + image = _load_image(image_file=image_file, gpu_image_decode=gpu_image_decode) + elif isinstance(image_file, str) and image_file.startswith("data:"): + image = _load_image(image_file=image_file, gpu_image_decode=gpu_image_decode) + elif isinstance(image_file, str): + image = _load_image(image_file=image_file, gpu_image_decode=gpu_image_decode) + else: + raise ValueError(f"Invalid image: {image_file}") + + return image, image_size + + +def get_image_bytes(image_file: str | bytes) -> bytes: + """Normalize various image inputs into raw bytes.""" + if isinstance(image_file, bytes): + return image_file + if image_file.startswith(("http://", "https://")): + timeout = int(os.getenv("REQUEST_TIMEOUT", "3")) + response = requests.get(image_file, timeout=timeout) + try: + response.raise_for_status() + result = response.content + finally: + response.close() + return result + if image_file.startswith("file://"): + with open(unquote(urlparse(image_file).path), "rb") as f: + return f.read() + if image_file.startswith("/"): + with open(image_file, "rb") as f: + return f.read() + if image_file.lower().endswith(image_extension_names): + with open(image_file, "rb") as f: + return f.read() + if isinstance(image_file, str) and image_file.startswith("data:"): + _, encoded = image_file.split(",", 1) + return pybase64.b64decode(encoded, validate=True) + if isinstance(image_file, str): + return pybase64.b64decode(image_file, validate=True) + raise NotImplementedError(f"Invalid image: {image_file}") + + +def load_audio( + audio_file: str | bytes, + sr: int | None = None, + mono: bool = True, +) -> np.ndarray: + # Use soundfile directly; librosa delegates to it and is moving away from + # audio loading support. + import soundfile as sf + from scipy.signal import resample + + if sr is None: + sr = 16000 + + if isinstance(audio_file, bytes): + audio, original_sr = sf.read(BytesIO(audio_file)) + elif audio_file.startswith("data:"): + _, encoded = audio_file.split(",", 1) + audio, original_sr = sf.read( + BytesIO(pybase64.b64decode(encoded, validate=True)) + ) + elif audio_file.startswith(("http://", "https://")): + timeout = int(os.getenv("REQUEST_TIMEOUT", "5")) + response = requests.get(audio_file, stream=True, timeout=timeout) + try: + response.raise_for_status() + audio, original_sr = sf.read(BytesIO(response.content)) + finally: + response.close() + elif isinstance(audio_file, str): + audio, original_sr = sf.read(audio_file) + else: + raise ValueError(f"Invalid audio format: {audio_file}") + + if original_sr != sr: + num_samples = int(len(audio) * float(sr) / original_sr) + audio = resample(audio, num_samples) + + if mono and len(audio.shape) > 1: + audio = np.mean(audio, axis=1) + + return audio + + +def set_ulimit(target_soft_limit=65535): + # number of open files + 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: + logger.warning("Failed to set RLIMIT_NOFILE: %s", e) + + # stack size + resource_type = resource.RLIMIT_STACK + current_soft, current_hard = resource.getrlimit(resource_type) + target_soft_limit_stack_size = 1024 * target_soft_limit + if current_soft < target_soft_limit_stack_size: + try: + resource.setrlimit( + resource_type, (target_soft_limit_stack_size, current_hard) + ) + except ValueError as e: + logger.warning("Failed to set RLIMIT_STACK: %s", e) + + +def add_api_key_middleware(app, api_key: str): + @app.middleware("http") + async def authentication(request, call_next): + if request.method == "OPTIONS": + return await call_next(request) + if request.url.path.startswith("/health"): + return await call_next(request) + if request.url.path.startswith("/metrics"): + return await call_next(request) + if request.headers.get("Authorization") != "Bearer " + api_key: + return ORJSONResponse(content={"error": "Unauthorized"}, status_code=401) + return await call_next(request) + + +def prepare_model_and_tokenizer(model_path: str, tokenizer_path: str): + from tokenspeed.runtime.utils.env import envs + + if envs.TOKENSPEED_USE_MODELSCOPE.get(): + if not os.path.exists(model_path): + from modelscope import snapshot_download + + model_path = snapshot_download(model_path) + tokenizer_path = snapshot_download( + tokenizer_path, ignore_patterns=["*.bin", "*.safetensors"] + ) + return model_path, tokenizer_path + + +def configure_logger(server_args, prefix: str = ""): + global LOG_PREFIX + LOG_PREFIX = prefix + + global LOG_LEVEL + LOG_LEVEL = server_args.log_level.upper() + + from tokenspeed._logging import suppress_noisy_third_party_logs + from tokenspeed.runtime.utils.env import envs + + suppress_noisy_third_party_logs() + + if TOKENSPEED_LOGGING_CONFIG_PATH := envs.TOKENSPEED_LOGGING_CONFIG_PATH.get(): + if not os.path.exists(TOKENSPEED_LOGGING_CONFIG_PATH): + raise FileNotFoundError( + "Setting TOKENSPEED_LOGGING_CONFIG_PATH from env with " + f"{TOKENSPEED_LOGGING_CONFIG_PATH} but it does not exist!" + ) + with open(TOKENSPEED_LOGGING_CONFIG_PATH, encoding="utf-8") as file: + custom_config = json.loads(file.read()) + logging.config.dictConfig(custom_config) + suppress_noisy_third_party_logs() + return + format = f"[%(asctime)s{prefix}] %(message)s" + log_level = getattr(logging, server_args.log_level.upper()) + logging.basicConfig( + level=log_level, + format=format, + datefmt="%Y-%m-%d %H:%M:%S", + force=True, + ) + + # Only set specified log level for tokenspeed-related loggers + for logger_name in logging.Logger.manager.loggerDict: + if "tokenspeed" in logger_name or logger_name.startswith("tokenspeed"): + logger_obj = logging.getLogger(logger_name) + if isinstance(logger_obj, logging.Logger): + logger_obj.setLevel(log_level) + for handler in logger_obj.handlers: + handler.setLevel(log_level) + + suppress_noisy_third_party_logs() + + +def set_weight_attrs( + weight: torch.Tensor, + weight_attrs: dict[str, Any] | None, +): + """Set attributes on a weight tensor. + + This method is used to set attributes on a weight tensor. This method + will not overwrite existing attributes. + + Args: + weight: The weight tensor. + weight_attrs: A dictionary of attributes to set on the weight tensor. + """ + if weight_attrs is None: + return + for key, value in weight_attrs.items(): + if hasattr(weight, key): + raise ValueError(f"Overwriting existing tensor attribute: {key}") + setattr(weight, key, value) + + +def broadcast_pyobj( + data: list[Any], + rank: int, + dist_group: torch.distributed.ProcessGroup | None = None, + src: int = 0, + force_cpu_device: bool = True, +): + """Broadcast inputs from src rank to all other ranks with torch.dist backend. + The `rank` here refer to the source rank on global process group (regardless + of dist_group argument). + """ + device = torch.device( + "cuda" if torch.cuda.is_available() and not force_cpu_device else "cpu" + ) + + if rank == src: + if len(data) == 0: + tensor_size = torch.tensor([0], dtype=torch.long, device=device) + dist.broadcast(tensor_size, src=src, group=dist_group) + else: + serialized_data = pickle.dumps(data) + size = len(serialized_data) + + tensor_data = torch.ByteTensor( + np.frombuffer(serialized_data, dtype=np.uint8) + ).to(device) + tensor_size = torch.tensor([size], dtype=torch.long, device=device) + + dist.broadcast(tensor_size, src=src, group=dist_group) + dist.broadcast(tensor_data, src=src, group=dist_group) + return data + else: + tensor_size = torch.tensor([0], dtype=torch.long, device=device) + dist.broadcast(tensor_size, src=src, group=dist_group) + size = tensor_size.item() + + if size == 0: + return [] + + tensor_data = torch.empty(size, dtype=torch.uint8, device=device) + dist.broadcast(tensor_data, src=src, group=dist_group) + + serialized_data = bytes(tensor_data.cpu().numpy()) + data = pickle.loads(serialized_data) + return data + + +step_counter = 0 + + +def get_zmq_socket( + context: zmq.Context, socket_type: zmq.SocketType, endpoint: str, bind: bool +) -> zmq.Socket: + mem = psutil.virtual_memory() + total_mem = mem.total / 1024**3 + available_mem = mem.available / 1024**3 + if total_mem > 32 and available_mem > 16: + buf_size = int(0.5 * 1024**3) + else: + buf_size = -1 + + socket = context.socket(socket_type) + if endpoint.find("[") != -1: + socket.setsockopt(zmq.IPV6, 1) + + def set_send_opt(): + socket.setsockopt(zmq.SNDHWM, 0) + socket.setsockopt(zmq.SNDBUF, buf_size) + + def set_recv_opt(): + socket.setsockopt(zmq.RCVHWM, 0) + socket.setsockopt(zmq.RCVBUF, buf_size) + + if socket_type == zmq.PUSH: + set_send_opt() + elif socket_type == zmq.PULL: + set_recv_opt() + elif socket_type == zmq.DEALER: + set_send_opt() + set_recv_opt() + else: + raise ValueError(f"Unsupported socket type: {socket_type}") + + if bind: + socket.bind(endpoint) + else: + socket.connect(endpoint) + + return socket + + +def delete_directory(dirpath): + try: + # This will remove the directory and all its contents + shutil.rmtree(dirpath) + except OSError as e: + logger.warning("Failed to delete directory %s: %s", dirpath, e.strerror) + + +# Temporary directory for prometheus multiprocess mode +# Cleaned up automatically when this object is garbage collected +prometheus_multiproc_dir: tempfile.TemporaryDirectory + + +def set_prometheus_multiproc_dir(): + # Set prometheus multiprocess directory + # tokenspeed uses prometheus multiprocess mode + # we need to set this before importing prometheus_client + # https://prometheus.github.io/client_python/multiprocess/ + global prometheus_multiproc_dir + + if "PROMETHEUS_MULTIPROC_DIR" in os.environ: + logger.debug("User set PROMETHEUS_MULTIPROC_DIR detected.") + prometheus_multiproc_dir = tempfile.TemporaryDirectory( + dir=os.environ["PROMETHEUS_MULTIPROC_DIR"] + ) + else: + prometheus_multiproc_dir = tempfile.TemporaryDirectory() + os.environ["PROMETHEUS_MULTIPROC_DIR"] = prometheus_multiproc_dir.name + logger.debug("PROMETHEUS_MULTIPROC_DIR: %s", os.environ["PROMETHEUS_MULTIPROC_DIR"]) + + +def add_prometheus_middleware(app): + # We need to import prometheus_client after setting the env variable `PROMETHEUS_MULTIPROC_DIR` + from prometheus_client import CollectorRegistry, make_asgi_app, multiprocess + + registry = CollectorRegistry() + multiprocess.MultiProcessCollector(registry) + metrics_route = Mount("/metrics", make_asgi_app(registry=registry)) + + # Workaround for 307 Redirect for /metrics + metrics_route.path_regex = re.compile("^/metrics(?P.*)$") + app.routes.append(metrics_route) + + +def get_amdgpu_memory_capacity(): + if not torch.cuda.is_available(): + raise RuntimeError( + "No AMD GPU available. Ensure ROCm drivers and a ROCm-enabled " + "PyTorch build are installed and accessible." + ) + + # Query each visible device's total memory (bytes) via the torch API + # (torch.cuda is reused for ROCm/HIP), and return the minimum in MiB so + # the value matches the previous rocminfo-based implementation. + memory_values = [ + torch.cuda.get_device_properties(i).total_memory // (1024 * 1024) + for i in range(torch.cuda.device_count()) + ] + + if not memory_values: + raise ValueError("No GPU memory values found.") + + return min(memory_values) + + +def get_nvgpu_memory_capacity(): + try: + # Run nvidia-smi and capture the output + result = subprocess.run( + ["nvidia-smi", "--query-gpu=memory.total", "--format=csv,noheader,nounits"], + capture_output=True, + text=True, + ) + + if result.returncode != 0: + raise RuntimeError(f"nvidia-smi error: {result.stderr.strip()}") + + # Parse the output to extract memory values + memory_values = [ + float(mem) + for mem in result.stdout.strip().split("\n") + if re.match(r"^\d+(\.\d+)?$", mem.strip()) + ] + + if not memory_values: + # Fallback to torch.cuda.mem_get_info() when failed to get memory capacity from nvidia-smi, + # typically in NVIDIA MIG mode. + if torch.cuda.is_available(): + logger.warning( + "Failed to get GPU memory capacity from nvidia-smi, falling back to torch.cuda.mem_get_info()." + ) + return torch.cuda.mem_get_info()[1] // 1024 // 1024 # unit: MB + raise ValueError("No GPU memory values found.") + + # Return the minimum memory value + return min(memory_values) + + except FileNotFoundError: + raise RuntimeError( + "nvidia-smi not found. Ensure NVIDIA drivers are installed and accessible." + ) + + +def crash_on_warnings(): + # Crash on warning if we are running CI tests + return get_bool_env_var("CI") or get_bool_env_var("GITHUB_ACTIONS") + + +def print_warning_once(msg: str) -> None: + # Set the stacklevel to 2 to print the caller's line info + logger.warning(msg, stacklevel=2) + + +def get_device_name(device_id: int = 0) -> str: + if hasattr(torch, "cuda") and torch.cuda.is_available(): + return torch.cuda.get_device_name(device_id) + + return "" + + +@lru_cache(maxsize=8) +def get_device(device_id: int | None = None) -> str: + if hasattr(torch, "cuda") and torch.cuda.is_available(): + if device_id is None: + return "cuda" + return f"cuda:{device_id}" + + raise RuntimeError("No accelerator (CUDA/ROCm) is available.") + + +def dataclass_to_string_truncated( + data, max_length=2048, skip_names: set[str] | None = None +): + if skip_names is None: + skip_names = set() + # Summarize tensors/ndarrays by shape — never str() the values (the bare + # str() fallthrough below would dump a whole multimodal feature tensor, + # bloating the request log). + if torch.is_tensor(data): + return f"Tensor(shape={tuple(data.shape)}, dtype={data.dtype})" + if isinstance(data, np.ndarray): + return f"ndarray(shape={tuple(data.shape)}, dtype={data.dtype})" + if isinstance(data, str): + if len(data) > max_length: + half_length = max_length // 2 + return f"{repr(data[:half_length])} ... {repr(data[-half_length:])}" + else: + return f"{repr(data)}" + elif isinstance(data, (list, tuple)): + # Recurse element-wise (was ``str(data)``, which would dump nested + # tensors in full) and propagate skip_names. + if len(data) > max_length: + half_length = max_length // 2 + shown = list(data[:half_length]) + ["..."] + list(data[-half_length:]) + else: + shown = data + inner = ", ".join( + ( + "..." + if x == "..." + else dataclass_to_string_truncated(x, max_length, skip_names) + ) + for x in shown + ) + return "[" + inner + "]" + elif isinstance(data, dict): + return ( + "{" + + ", ".join( + f"'{k}': {dataclass_to_string_truncated(v, max_length, skip_names)}" + for k, v in data.items() + if k not in skip_names + ) + + "}" + ) + elif dataclasses.is_dataclass(data): + fields = dataclasses.fields(data) + return ( + f"{data.__class__.__name__}(" + + ", ".join( + f"{f.name}={dataclass_to_string_truncated(getattr(data, f.name), max_length, skip_names)}" + for f in fields + if f.name not in skip_names + ) + + ")" + ) + else: + return str(data) + + +class MultiprocessingSerializer: + @staticmethod + def serialize(obj, output_str: bool = False): + """ + Serialize a Python object using ForkingPickler. + + Args: + obj: The object to serialize. + output_str (bool): If True, return a base64-encoded string instead of raw bytes. + + Returns: + bytes or str: The serialized object. + """ + buf = io.BytesIO() + ForkingPickler(buf).dump(obj) + buf.seek(0) + output = buf.read() + + if output_str: + # Convert bytes to base64-encoded string + output = pybase64.b64encode(output).decode("utf-8") + + return output + + @staticmethod + def deserialize(data): + """ + Deserialize a previously serialized object. + + Args: + data (bytes or str): The serialized data, optionally base64-encoded. + + Returns: + The deserialized Python object. + """ + if isinstance(data, str): + # Decode base64 string to bytes + data = pybase64.b64decode(data, validate=True) + + return ForkingPickler.loads(data) + + +def debug_timing(func): + def wrapper(*args, **kwargs): + if logger.isEnabledFor(logging.DEBUG): + tic = torch.cuda.Event(enable_timing=True) + toc = torch.cuda.Event(enable_timing=True) + tic.record() + result = func(*args, **kwargs) + toc.record() + toc.synchronize() # Wait for the function to complete without synchronizing all ops on the GPU + elapsed = tic.elapsed_time(toc) + indices = kwargs.get("indices", args[1] if len(args) > 1 else None) + num_tokens = len(indices) if indices is not None else 0 + throughput = num_tokens / elapsed * 1000 if elapsed > 0 else 0 + logger.debug( + "Transfer time: %s ms, throughput: %s tokens/s", elapsed, throughput + ) + return result + else: + return func(*args, **kwargs) + + return wrapper + + +def nullable_str(val: str): + if not val or val == "None": + return None + return val + + +def is_valid_ipv6_address(address: str) -> bool: + try: + ipaddress.IPv6Address(address) + return True + except ValueError: + return False + + +def launch_dummy_health_check_server(host, port, enable_metrics): + + import uvicorn + from fastapi import FastAPI, Response + + app = FastAPI() + + @app.get("/health") + async def health(): + """Check the health of the http server.""" + return Response(status_code=200) + + @app.get("/health_generate") + async def health_generate(): + """Check the health of the http server.""" + return Response(status_code=200) + + # Add prometheus middleware + if enable_metrics: + add_prometheus_middleware(app) + enable_func_timer() + + config = uvicorn.Config( + app, + host=host, + port=port, + timeout_keep_alive=5, + loop="auto", + log_config=None, + log_level="warning", + ) + server = uvicorn.Server(config=config) + + try: + loop = asyncio.get_running_loop() + logger.info( + "Dummy health check server scheduled on existing loop at %s:%s", host, port + ) + loop.create_task(server.serve()) + + except RuntimeError: + logger.info("Starting dummy health check server at %s:%s", host, port) + server.run() + + +def set_cuda_arch(): + platform = current_platform() + if not platform.is_nvidia: + return + + arch = f"{platform.arch_version.major}.{platform.arch_version.minor}" + os.environ["TORCH_CUDA_ARCH_LIST"] = f"{arch}{'+PTX' if arch == '9.0' else ''}" + + +def next_power_of_2(n: int): + return 1 << (n - 1).bit_length() if n > 0 else 1 + + +def round_up(x: int, y: int) -> int: + return ((x - 1) // y + 1) * y + + +setattr(triton, "next_power_of_2", next_power_of_2) + + +def add_prefix(name: str, prefix: str) -> str: + """Add a weight path prefix to a module name. + + Args: + name: base module name. + prefix: weight prefix str to added to the front of `name` concatenated with `.`. + + Returns: + The string `prefix.name` if prefix is non-empty, otherwise just `name`. + """ + return name if not prefix else f"{prefix}.{name}" + + +# Can be more general if it is used in multiple places (keep it simple and thus not general now) + + +def log_info_on_rank0(logger, msg): + import torch.distributed as dist + + if not dist.is_initialized() or dist.get_rank() == 0: + logger.info(msg) + + +T = TypeVar("T") + + +class Withable(Generic[T]): + def __init__(self): + self._value: T | None = None + + @property + def value(self) -> T: + return self._value + + @contextmanager + def with_value(self, new_value: T): + if self._value is not None: + raise RuntimeError("Withable value is already set.") + self._value = new_value + try: + yield + finally: + if self._value is not new_value: + raise RuntimeError("Withable value changed while context was active.") + self._value = None + + +def find_local_repo_dir(repo_id: str, revision: str | None = None) -> str | None: + import huggingface_hub as hf + + # Build cache path + cache_path = os.path.join( + hf.constants.HF_HUB_CACHE, + hf.constants.REPO_ID_SEPARATOR.join(["models", *repo_id.split("/")]), + ) + + # Get revision from main ref if not specified + if not revision: + ref_path = os.path.join(cache_path, "refs", "main") + if os.path.isfile(ref_path): + with open(ref_path) as f: + revision = f.read().strip() + + # List files from revision directory + if revision: + rev_dir = os.path.join(cache_path, "snapshots", revision) + if os.path.isdir(rev_dir): + return rev_dir + + return None + + +def read_system_prompt_from_file(model_name: str) -> str: + """Read system prompt from a file in the HuggingFace cache directory. + + Args: + model_name: The model name to construct the file path + + Returns: + The system prompt content from the file, or empty string if file not found + """ + try: + local_repo_dir = find_local_repo_dir(model_name) + if local_repo_dir: + system_prompt_file = os.path.join(local_repo_dir, "SYSTEM_PROMPT.txt") + if os.path.exists(system_prompt_file): + with open(system_prompt_file, encoding="utf-8") as f: + return f.read() + + return "" + except Exception: + # If anything fails, return empty string + return "" + + +class LazyValue: + def __init__(self, creator: Callable): + self._creator = creator + self._value = None + + @property + def value(self): + if self._creator is not None: + self._value = self._creator() + self._creator = None + return self._value + + +def ceil_div(x: int, y: int) -> int: + return (x + y - 1) // y + + +# Only physical cores are used. Logical cores are excluded. + + +def lru_cache_frozenset(maxsize=128): + def _to_hashable(o): + try: + hash(o) + return o + except TypeError: + # Not hashable; convert based on type + if isinstance(o, (dict)): + return frozenset( + (_to_hashable(k), _to_hashable(v)) for k, v in o.items() + ) + elif isinstance(o, set): + return frozenset(_to_hashable(v) for v in o) + elif isinstance(o, (list, tuple)) or ( + isinstance(o, Sequence) and not isinstance(o, (str, bytes)) + ): + return tuple(_to_hashable(v) for v in o) + else: + raise TypeError(f"Cannot make hashable: {type(o)}") + + def decorator(func): + cache = OrderedDict() + + @functools.wraps(func) + def wrapper(*args, **kwargs): + h_args = tuple(_to_hashable(a) for a in args) + h_kwargs = frozenset( + (_to_hashable(k), _to_hashable(v)) for k, v in kwargs.items() + ) + key = (h_args, h_kwargs) + if key in cache: + cache.move_to_end(key) + return cache[key] + result = func(*args, **kwargs) + cache[key] = result + if maxsize is not None and len(cache) > maxsize: + cache.popitem(last=False) + return result + + wrapper.cache_clear = cache.clear # For manual cache clearing + return wrapper + + return decorator + + +LOG_PREFIX = None +LOG_LEVEL = "INFO" + + +class CustomFormatter(logging.Formatter): + grey = "\x1b[38;20m" + yellow = "\x1b[33;20m" + red = "\x1b[31;20m" + bold_red = "\x1b[31;1m" + reset = "\x1b[0m" + + FORMATS = None + + def format(self, record): + if self.FORMATS is None: + format = f"[%(asctime)s {LOG_PREFIX}] - %(levelname)s - %(message)s (%(filename)s:%(lineno)d)" + self.FORMATS = { + logging.DEBUG: self.grey + format + self.reset, + logging.INFO: self.grey + format + self.reset, + logging.WARNING: self.yellow + format + self.reset, + logging.ERROR: self.red + format + self.reset, + logging.CRITICAL: self.bold_red + format + self.reset, + } + + log_fmt = self.FORMATS.get(record.levelno) + formatter = logging.Formatter(log_fmt) + return formatter.format(record) + + +def get_colorful_logger(name): + logger = logging.getLogger(name) + logger.propagate = False + logger.setLevel(LOG_LEVEL) + + ch = logging.StreamHandler() + ch.setLevel(LOG_LEVEL) + ch.setFormatter(CustomFormatter()) + # ch.flush = lambda: True + + logger.addHandler(ch) + logger.propagate = False + return logger + + +def _maybe_json_dict(path: str | os.PathLike) -> dict[str, str]: + with open(path) as f: + try: + return json.loads(f.read()) + except Exception: + return dict[str, str]() + + +def _maybe_space_split_dict(path: str | os.PathLike) -> dict[str, str]: + parsed_dict = dict[str, str]() + with open(path) as f: + for line in f.readlines(): + try: + model_name, redirect_name = line.strip().split() + parsed_dict[model_name] = redirect_name + except Exception: + pass + return parsed_dict + + +def maybe_model_redirect(model: str) -> str: + """ + Use model_redirect to redirect the model name to a local folder. + + :param model: hf model name + :return: maybe redirect to a local folder + """ + + from tokenspeed.runtime.utils.env import envs + + model_redirect_path = envs.TOKENSPEED_MODEL_REDIRECT_PATH.get() + + if not model_redirect_path: + return model + + if not Path(model_redirect_path).exists(): + return model + + redirect_dict = _maybe_json_dict(model_redirect_path) or _maybe_space_split_dict( + model_redirect_path + ) + if redirect_model := redirect_dict.get(model): + logger.info("model redirect: [ %s ] -> [ %s ]", model, redirect_model) + return redirect_model + + return model + + +def random_uuid() -> str: + return str(uuid.uuid4().hex) + + +def flatten_nested_list(nested_list): + if isinstance(nested_list, list): + return [ + item for sublist in nested_list for item in flatten_nested_list(sublist) + ] + else: + return [nested_list] + + +def convert_json_schema_to_str(json_schema: dict | str | type[BaseModel]) -> str: + """Convert a JSON schema to a string. + Parameters + ---------- + json_schema + The JSON schema. + Returns + ------- + str + The JSON schema converted to a string. + Raises + ------ + ValueError + If the schema is not a dictionary, a string or a Pydantic class. + """ + if isinstance(json_schema, dict): + schema_str = json.dumps(json_schema) + elif isinstance(json_schema, str): + schema_str = json_schema + elif issubclass(json_schema, BaseModel): + schema_str = json.dumps(json_schema.model_json_schema()) + else: + raise ValueError( + f"Cannot parse schema {json_schema}. The schema must be either " + + "a Pydantic class, a dictionary or a string that contains the JSON " + + "schema specification" + ) + return schema_str diff --git a/python/tokenspeed/runtime/utils/cuda_stream.py b/python/tokenspeed/runtime/utils/cuda_stream.py new file mode 100644 index 0000000..17191d9 --- /dev/null +++ b/python/tokenspeed/runtime/utils/cuda_stream.py @@ -0,0 +1,56 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from contextlib import contextmanager + +import torch + + +class StreamFork: + def __init__(self, aux_stream: torch.cuda.Stream | None): + self.aux_stream = aux_stream + self.fork_event = torch.cuda.Event() if aux_stream is not None else None + self.join_event = torch.cuda.Event() if aux_stream is not None else None + self._active = False + self._current: torch.cuda.Stream | None = None + + @contextmanager + def scope(self, *, enable: bool): + self._active = enable and self.aux_stream is not None + if self._active: + self._current = torch.cuda.current_stream() + self.fork_event.record(self._current) + try: + yield self + finally: + if self._active: + self.join_event.wait(self._current) + self._active = False + self._current = None + + @contextmanager + def branch(self): + if not self._active: + yield + return + with torch.cuda.stream(self.aux_stream): + self.fork_event.wait(self.aux_stream) + yield + self.join_event.record(self.aux_stream) diff --git a/python/tokenspeed/runtime/utils/custom_ops.py b/python/tokenspeed/runtime/utils/custom_ops.py new file mode 100644 index 0000000..1a9fb2d --- /dev/null +++ b/python/tokenspeed/runtime/utils/custom_ops.py @@ -0,0 +1,61 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from __future__ import annotations + +import importlib +from collections.abc import Callable + +import torch +import torch.library +from torch.library import Library + +tokenspeed_lib = Library("tokenspeed", "FRAGMENT") + + +def direct_register_custom_op( + op_name: str, + op_func: Callable, + mutates_args: list[str], + fake_impl: Callable | None = None, + target_lib: Library | None = None, +) -> None: + """Register a low-overhead torch custom op in the TokenSpeed namespace.""" + + target = target_lib or tokenspeed_lib + lib_name = getattr(getattr(target, "m", None), "name", "tokenspeed") + try: + if hasattr(torch.ops, lib_name) and hasattr( + getattr(torch.ops, lib_name), op_name + ): + return + except (AttributeError, RuntimeError): + pass + + if hasattr(torch.library, "infer_schema"): + schema_str = torch.library.infer_schema(op_func, mutates_args=mutates_args) + else: + custom_op_impl = importlib.import_module("torch._custom_op.impl") + schema_str = custom_op_impl.infer_schema(op_func, mutates_args) + + target.define(op_name + schema_str) + target.impl(op_name, op_func, "CUDA") + if fake_impl is not None: + target._register_fake(op_name, fake_impl) diff --git a/python/tokenspeed/runtime/utils/dispatch.py b/python/tokenspeed/runtime/utils/dispatch.py new file mode 100644 index 0000000..c6ea0a4 --- /dev/null +++ b/python/tokenspeed/runtime/utils/dispatch.py @@ -0,0 +1,44 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from collections.abc import Callable +from typing import Any + + +class TypeBasedDispatcher: + def __init__(self, mapping: list[tuple[type, Callable]]): + self._mapping = mapping + self._fallback_fn = None + + def add_fallback_fn(self, fallback_fn: Callable): + self._fallback_fn = fallback_fn + + def __iadd__(self, other: "TypeBasedDispatcher"): + self._mapping.extend(other._mapping) + return self + + def __call__(self, obj: Any): + for ty, fn in self._mapping: + if isinstance(obj, ty): + return fn(obj) + + if self._fallback_fn is not None: + return self._fallback_fn(obj) + raise ValueError(f"Invalid object: {obj}") diff --git a/python/tokenspeed/runtime/utils/env.py b/python/tokenspeed/runtime/utils/env.py new file mode 100755 index 0000000..540ddb5 --- /dev/null +++ b/python/tokenspeed/runtime/utils/env.py @@ -0,0 +1,304 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import os +import warnings +from contextlib import contextmanager +from typing import Any + +from tokenspeed.runtime.utils.pdl import pdl_enabled +from tokenspeed.runtime.utils.server_args import ServerArgs + +global_server_args_dict: dict = { + "attention_backend": ServerArgs.attention_backend, + "sampling_backend": ServerArgs.sampling_backend, + "attention_use_fp4_indexer_cache": ServerArgs.attention_use_fp4_indexer_cache, + "deepseek_v4_mega_moe_max_num_tokens": ServerArgs.deepseek_v4_mega_moe_max_num_tokens, + "deepseek_v4_indexer_prefill_max_logits_mb": ServerArgs.deepseek_v4_indexer_prefill_max_logits_mb, + "deepseek_v4_prefill_chunk_size": ServerArgs.deepseek_v4_prefill_chunk_size, + "triton_attention_reduce_in_fp32": ServerArgs.triton_attention_reduce_in_fp32, + "kv_cache_dtype": ServerArgs.kv_cache_dtype, + "enable_nan_detection": ServerArgs.enable_nan_detection, + "enable_p2p_check": ServerArgs.enable_p2p_check, + "mapping": ServerArgs.mapping, + "force_deterministic_rsag": ServerArgs.force_deterministic_rsag, + "low_latency_max_num_tokens_per_gpu": ServerArgs.low_latency_max_num_tokens_per_gpu, + "device": ServerArgs.device, + "draft_model_path_use_base": ServerArgs.draft_model_path_use_base, + "disable_pdl": ServerArgs.disable_pdl, + "enable_prefix_caching": ServerArgs.enable_prefix_caching, + "mla_disable_ragged": ServerArgs.mla_disable_ragged, + "chunked_prefill_size": ServerArgs.chunked_prefill_size, + "mla_chunk_multiplier": ServerArgs.mla_chunk_multiplier, + "ep_num_redundant_experts": ServerArgs.ep_num_redundant_experts, + "ep_dispatch_algorithm": ServerArgs.ep_dispatch_algorithm, + "enable_eplb": ServerArgs.enable_eplb, + "mm_attention_backend": ServerArgs.mm_attention_backend, + "comm_fusion_max_num_tokens": ServerArgs.comm_fusion_max_num_tokens, + "enable_allreduce_fusion": ServerArgs.enable_allreduce_fusion, + "max_prefill_tokens": ServerArgs.max_prefill_tokens, + "max_model_len": ServerArgs.max_model_len, + "max_num_seqs": ServerArgs.max_num_seqs, + "moe_backend": ServerArgs.moe_backend, + "enforce_eager": ServerArgs.enforce_eager, + "max_cudagraph_capture_size": ServerArgs.max_cudagraph_capture_size, + "cudagraph_capture_sizes": ServerArgs.cudagraph_capture_sizes, + "disable_prefill_graph": ServerArgs.disable_prefill_graph, + "prefill_graph_max_tokens": ServerArgs.prefill_graph_max_tokens, + "mamba_track_interval": ServerArgs.mamba_track_interval, + "all2all_backend": ServerArgs.all2all_backend, +} + + +def global_server_args_dict_update(server_args: ServerArgs): + global_server_args_dict.update( + { + "attention_backend": server_args.attention_backend, + "sampling_backend": server_args.sampling_backend, + "attention_use_fp4_indexer_cache": server_args.attention_use_fp4_indexer_cache, + "deepseek_v4_mega_moe_max_num_tokens": server_args.deepseek_v4_mega_moe_max_num_tokens, + "deepseek_v4_indexer_prefill_max_logits_mb": server_args.deepseek_v4_indexer_prefill_max_logits_mb, + "deepseek_v4_prefill_chunk_size": server_args.deepseek_v4_prefill_chunk_size, + "triton_attention_reduce_in_fp32": server_args.triton_attention_reduce_in_fp32, + "kv_cache_dtype": server_args.kv_cache_dtype, + "enable_nan_detection": server_args.enable_nan_detection, + "enable_p2p_check": server_args.enable_p2p_check, + "mapping": server_args.mapping, + "force_deterministic_rsag": server_args.force_deterministic_rsag, + "low_latency_max_num_tokens_per_gpu": server_args.low_latency_max_num_tokens_per_gpu, + "device": server_args.device, + "draft_model_path_use_base": server_args.draft_model_path_use_base, + "speculative_algorithm": server_args.speculative_algorithm, + "speculative_num_draft_tokens": server_args.speculative_num_draft_tokens, + "disable_pdl": server_args.disable_pdl, + "enable_prefix_caching": server_args.enable_prefix_caching, + "mla_disable_ragged": server_args.mla_disable_ragged, + "chunked_prefill_size": server_args.chunked_prefill_size, + "mla_chunk_multiplier": server_args.mla_chunk_multiplier, + "ep_num_redundant_experts": server_args.ep_num_redundant_experts, + "ep_dispatch_algorithm": server_args.ep_dispatch_algorithm, + "enable_eplb": server_args.enable_eplb, + "mm_attention_backend": server_args.mm_attention_backend, + "comm_fusion_max_num_tokens": server_args.comm_fusion_max_num_tokens, + "enable_allreduce_fusion": server_args.enable_allreduce_fusion, + "max_prefill_tokens": server_args.max_prefill_tokens, + "max_model_len": server_args.max_model_len, + "max_num_seqs": server_args.max_num_seqs, + "moe_backend": server_args.moe_backend, + "enforce_eager": server_args.enforce_eager, + "max_cudagraph_capture_size": server_args.max_cudagraph_capture_size, + "cudagraph_capture_sizes": server_args.cudagraph_capture_sizes, + "disable_prefill_graph": server_args.disable_prefill_graph, + "prefill_graph_max_tokens": server_args.prefill_graph_max_tokens, + "all2all_backend": server_args.all2all_backend, + } + ) + pdl_enabled.cache_clear() + + +class EnvField: + def __init__(self, default: Any): + self.default = default + # we use None to indicate whether the value is set or not + # If the value is manually set to None, we need mark it as _set_to_none. + # Always use clear() to reset the value, which leads to the default fallback. + self._set_to_none = False + + def __set_name__(self, owner, name): + self.name = name + + def parse(self, value: str) -> Any: + raise NotImplementedError() + + def get(self) -> Any: + value = os.getenv(self.name) + if self._set_to_none: + if value is not None: + raise RuntimeError(f"{self.name} is set while marked as None.") + return None + + if value is None: + return self.default + + try: + return self.parse(value) + except ValueError as e: + warnings.warn( + f'Invalid value for {self.name}: {e}, using default "{self.default}"' + ) + return self.default + + def is_set(self): + # If None is manually set, it is considered as set. + return self.name in os.environ or self._set_to_none + + def get_set_value_or(self, or_value: Any): + # Ugly usage, but only way to get custom default value. + return self.get() if self.is_set() else or_value + + def set(self, value: Any): + if value is None: + self._set_to_none = True + os.environ.pop(self.name, None) + else: + self._set_to_none = False + os.environ[self.name] = str(value) + + @contextmanager + def override(self, value: Any): + backup_present = self.name in os.environ + backup_value = os.environ.get(self.name) + backup_set_to_none = self._set_to_none + self.set(value) + yield + if backup_present: + os.environ[self.name] = backup_value + else: + os.environ.pop(self.name, None) + self._set_to_none = backup_set_to_none + + def clear(self): + os.environ.pop(self.name, None) + self._set_to_none = False + + @property + def value(self): + return self.get() + + def __bool__(self): + raise RuntimeError( + "Please use `envs.YOUR_FLAG.get()` instead of `envs.YOUR_FLAG`" + ) + + def __len__(self): + raise RuntimeError( + "Please use `envs.YOUR_FLAG.get()` instead of `envs.YOUR_FLAG`" + ) + + +class EnvStr(EnvField): + def parse(self, value: str) -> str: + return value + + +class EnvBool(EnvField): + def parse(self, value: str) -> bool: + value = value.lower() + if value in ["true", "1", "yes", "y"]: + return True + if value in ["false", "0", "no", "n"]: + return False + raise ValueError(f'"{value}" is not a valid boolean value') + + +class EnvInt(EnvField): + def parse(self, value: str) -> int: + try: + return int(value) + except ValueError: + raise ValueError(f'"{value}" is not a valid integer value') + + +class EnvFloat(EnvField): + def parse(self, value: str) -> float: + try: + return float(value) + except ValueError: + raise ValueError(f'"{value}" is not a valid float value') + + +class Envs: + # fmt: off + + # Model download + TOKENSPEED_USE_MODELSCOPE = EnvBool(False) + + # Test and debug + TOKENSPEED_CUDA_COREDUMP = EnvBool(False) + TOKENSPEED_CUDA_COREDUMP_DIR = EnvStr("/tmp/tokenspeed_cuda_coredumps") + TOKENSPEED_PROFILE_WITH_STACK = EnvBool(True) + TOKENSPEED_TEST_REQUEST_TIME_STATS = EnvBool(False) + TOKENSPEED_PROFILER_DIR = EnvStr("/tmp") + TOKENSPEED_CI_SMALL_KV_SIZE = EnvInt(-1) + TOKENSPEED_NVTX = EnvBool(False) + TOKENSPEED_DP_SAMPLING_BACKEND = EnvStr(None) + + # Scheduler + TOKENSPEED_BLOCK_NONZERO_RANK_CHILDREN = EnvBool(True) + + # Mooncake + TOKENSPEED_KVSTORE_MOONCAKE_CONFIG_PATH = EnvStr(None) + MOONCAKE_MASTER = EnvStr(None) + MOONCAKE_LOCAL_HOSTNAME = EnvStr("localhost") + MOONCAKE_TE_META_DATA_SERVER = EnvStr("P2PHANDSHAKE") + MOONCAKE_GLOBAL_SEGMENT_SIZE = EnvStr("4gb") + MOONCAKE_PROTOCOL = EnvStr("tcp") + MOONCAKE_DEVICE = EnvStr("") + MOONCAKE_MASTER_METRICS_PORT = EnvInt(9003) + MOONCAKE_CHECK_SERVER = EnvBool(False) + TOKENSPEED_DISAGGREGATION_FAILED_SESSION_TTL = EnvInt(30) + TOKENSPEED_DISAGGREGATION_HEARTBEAT_INTERVAL = EnvFloat(5.0) + TOKENSPEED_DISAGGREGATION_HEARTBEAT_MAX_FAILURE = EnvInt(2) + TOKENSPEED_DISAGGREGATION_QUEUE_SIZE = EnvInt(4) + TOKENSPEED_DISAGGREGATION_THREAD_POOL_SIZE = EnvInt(-1) + TOKENSPEED_DISAGGREGATION_BOOTSTRAP_TIMEOUT = EnvInt(120) + TOKENSPEED_DISAGGREGATION_WAITING_TIMEOUT = EnvInt(300) + TOKENSPEED_EPD_ENCODE_RING_SLOTS = EnvInt(64) + TOKENSPEED_EPD_ENCODE_RING_SLOT_MB = EnvInt(None) + TOKENSPEED_EPD_ENCODE_EMBED_CACHE_MB = EnvInt(4096) + TOKENSPEED_EPD_ENCODE_EMBED_CACHE_DRAM_MB = EnvInt(0) + TOKENSPEED_EPD_RECV_POOL_SLOTS = EnvInt(16) + TOKENSPEED_EPD_RECV_POOL_SLOT_MB = EnvInt(256) + TOKENSPEED_EPD_EMBEDDING_SHARD = EnvBool(True) + TOKENSPEED_PD_LAYERWISE_DEBUG = EnvBool(False) + TOKENSPEED_PD_PREFILL_METADATA_TIMEOUT = EnvFloat(5.0) + + # Quantization + TOKENSPEED_NVFP4_GEMM_SWIGLU_NVFP4_QUANT = EnvBool(True) + TOKENSPEED_MINIMAX_AR_USE_TRITON = EnvBool(False) + + # EPLB + TOKENSPEED_EXPERT_DISTRIBUTION_RECORDER_DIR = EnvStr("/tmp") + + # Runtime behavior + TOKENSPEED_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN = EnvBool(False) + TOKENSPEED_DETOKENIZER_MAX_STATES = EnvInt(1 << 16) + TOKENSPEED_FORCE_FAKE_FULL_NVLINK = EnvBool(False) + TOKENSPEED_HEALTH_CHECK_TIMEOUT = EnvInt(20) + TOKENSPEED_HOST_IP = EnvStr("") + TOKENSPEED_LOGGING_CONFIG_PATH = EnvStr(None) + TOKENSPEED_MAMBA_SSM_DTYPE = EnvStr("float32") + TOKENSPEED_MODEL_REDIRECT_PATH = EnvStr(None) + TOKENSPEED_MOE_PADDING = EnvBool(False) + TOKENSPEED_MOE_CONFIG_DIR = EnvStr(None) + TOKENSPEED_ENABLE_TORCH_INFERENCE_MODE = EnvBool(True) + TOKENSPEED_NUMA_AWARE_WORKER_AFFINITY = EnvBool(True) + TOKENSPEED_REQUEST_CONVERSION_WORKERS = EnvInt(8) + + # Multimodal / VLM + TOKENSPEED_LOG_MM_TIMING = EnvBool(False) + TOKENSPEED_MM_ENABLE_ENCODER_CUDA_GRAPH = EnvBool(False) + TOKENSPEED_MM_VIDEO_ENCODER_CUDA_GRAPH_MAX_SEQUENCES_PER_BATCH = EnvInt(None) + TOKENSPEED_MM_SKIP_COMPUTE_HASH = EnvBool(False) + + # fmt: on + + +envs = Envs() diff --git a/python/tokenspeed/runtime/utils/exceptions.py b/python/tokenspeed/runtime/utils/exceptions.py new file mode 100644 index 0000000..3b233b3 --- /dev/null +++ b/python/tokenspeed/runtime/utils/exceptions.py @@ -0,0 +1,27 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import sys +import traceback + + +def get_exception_traceback(): + etype, value, tb = sys.exc_info() + return "".join(traceback.format_exception(etype, value, tb)) diff --git a/python/tokenspeed/runtime/utils/flashinfer_config.py b/python/tokenspeed/runtime/utils/flashinfer_config.py new file mode 100644 index 0000000..f1928b8 --- /dev/null +++ b/python/tokenspeed/runtime/utils/flashinfer_config.py @@ -0,0 +1,29 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import os + +DEFAULT_FLASHINFER_WORKSPACE_SIZE = 384 * 1024 * 1024 + + +def get_flashinfer_workspace_size() -> int: + return int( + os.environ.get("FLASHINFER_WORKSPACE_SIZE", DEFAULT_FLASHINFER_WORKSPACE_SIZE) + ) diff --git a/python/tokenspeed/runtime/utils/hf_transformers_utils.py b/python/tokenspeed/runtime/utils/hf_transformers_utils.py new file mode 100755 index 0000000..0f18382 --- /dev/null +++ b/python/tokenspeed/runtime/utils/hf_transformers_utils.py @@ -0,0 +1,617 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Utilities for Huggingface Transformers.""" + +import contextlib +import copy +import importlib.util +import json +import logging +import os +import warnings +from collections.abc import Callable +from typing import Any + +import torch +from huggingface_hub import snapshot_download +from transformers import ( + AutoConfig, + AutoTokenizer, + GenerationConfig, + PretrainedConfig, + PreTrainedTokenizer, + PreTrainedTokenizerFast, +) +from transformers.utils import cached_file + +from tokenspeed.runtime.configs import ( + DeepseekV4Config, + KimiK2Config, + KimiK25Config, + MiniMaxM2Config, + Qwen2Config, + Qwen3_5Config, + Qwen3_5MoeConfig, + Qwen3ASRConfig, + Qwen3Config, + Qwen3MoeConfig, +) +from tokenspeed.runtime.utils import lru_cache_frozenset + +_CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = { + Qwen2Config.model_type: Qwen2Config, + Qwen3Config.model_type: Qwen3Config, + Qwen3MoeConfig.model_type: Qwen3MoeConfig, + Qwen3ASRConfig.model_type: Qwen3ASRConfig, + DeepseekV4Config.model_type: DeepseekV4Config, + Qwen3_5Config.model_type: Qwen3_5Config, + Qwen3_5MoeConfig.model_type: Qwen3_5MoeConfig, + MiniMaxM2Config.model_type: MiniMaxM2Config, + KimiK2Config.model_type: KimiK2Config, + KimiK25Config.model_type: KimiK25Config, +} + +_DEEPSEEK_V4_ENCODING_MODULE_NAME = "_tokenspeed_deepseek_v4_encoding" + +for name, cls in _CONFIG_REGISTRY.items(): + with contextlib.suppress(ValueError): + AutoConfig.register(name, cls) + + +def resolve_architecture(config: PretrainedConfig) -> str: + """Return ``config.architectures[0]`` or the config class name. + + ``config.architectures`` can be ``None`` on configs that forward + attribute access to a nested ``text_config`` (e.g. ``Qwen3_5MoeConfig``). + Callers should use this helper instead of indexing the list directly. + """ + archs = getattr(config, "architectures", None) + if archs: + return archs[0] + return type(config).__name__ + + +def get_hf_text_config(config: PretrainedConfig): + """Get the "sub" config relevant to llm for multi modal models. + No op for pure text models. + """ + class_name = resolve_architecture(config) + 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. + # We set `dtype` of config to `torch.float16` for the weights, as + # `torch.float16` is default used for image features in + # `python/tokenspeed/runtime/models/llava.py`. + config.dtype = torch.float16 + return config + + text_config = None + if hasattr(config, "text_config"): + # The code operates under the assumption that text_config should have + # `num_attention_heads` (among others). Check here to fail early + # if transformers config doesn't align with this assumption. + if not hasattr(config.text_config, "num_attention_heads"): + raise AttributeError("text_config must define num_attention_heads.") + text_config = config.text_config + if hasattr(config, "language_config"): + text_config = config.language_config + if hasattr(config, "thinker_config"): + # Qwen Omni wrappers keep the language model below thinker_config. + thinker_config = config.thinker_config + if hasattr(thinker_config, "text_config"): + thinker_config.text_config.dtype = thinker_config.dtype + text_config = thinker_config.text_config + else: + text_config = thinker_config + + if text_config is None: + return config + + if hasattr(config, "quantization_config") and not hasattr( + text_config, "quantization_config" + ): + quantization_config = config.quantization_config + for key in ["ignore", "ignored_layers", "modules_to_not_convert"]: + if key in quantization_config and isinstance( + quantization_config[key], list + ): + quantization_config[key] = [ + ( + x.replace("language_model.", "") + if x.startswith("language_model.") + else x + ) + for x in quantization_config[key] + ] + text_config.quantization_config = quantization_config + + return text_config + + +def _materialize_architectures(config: PretrainedConfig, raw_config: dict) -> None: + """Ensure ``config.architectures`` resolves to a real ``list[str]``. + + HuggingFace's ``from_pretrained`` sometimes returns a config whose + ``.architectures`` attribute resolves to ``None`` via ``__getattr__`` + forwarding to a nested text_config (observed on ``Qwen3_5MoeConfig``; + likely to repeat on any wrapper class with the same pattern). The + on-disk ``config.json`` is the source of truth, so pin its value + onto ``config.__dict__`` when the live config has lost it. Bypasses + ``__setattr__`` deliberately — that's the only way around the + ``__getattr__`` redirect. + + Silently no-ops when the raw value is missing, empty, or not a + ``list[str]``; downstream code already handles the absence via + ``resolve_architecture``. + """ + if getattr(config, "architectures", None): + return + raw_archs = raw_config.get("architectures") + if not ( + isinstance(raw_archs, list) + and raw_archs + and all(isinstance(a, str) for a in raw_archs) + ): + return + config.__dict__["architectures"] = list(raw_archs) + + +def _restore_raw_glm_dsa_fields(config: PretrainedConfig, raw_config: dict) -> None: + if raw_config.get("architectures") != ["GlmMoeDsaForCausalLM"]: + return + + # Transformers may rewrite these GLM DSA dimensions; config.json is authoritative. + for key in ( + "qk_head_dim", + "qk_nope_head_dim", + "qk_rope_head_dim", + "v_head_dim", + "kv_lora_rank", + "q_lora_rank", + "index_topk", + "index_head_dim", + "index_n_heads", + "index_topk_freq", + "index_skip_topk_offset", + "index_topk_pattern", + "indexer_types", + "indexer_rope_interleave", + "index_share_for_mtp_iteration", + ): + if key in raw_config: + setattr(config, key, raw_config[key]) + + +def get_config( + model: str, + trust_remote_code: bool, + revision: str | None = None, + model_override_args: dict | None = None, + is_draft_worker: bool | None = False, + **kwargs, +): + if os.path.isdir(model): + model_path = model + else: + model_path = snapshot_download( + model, ignore_patterns=["*.pt", "*.safetensors", "*.bin"] + ) + + try: + with open(os.path.join(model_path, "config.json")) as file: + raw_config = json.load(file) + except FileNotFoundError: + raise RuntimeError(f"Config file not found in {model}. Please check the path.") + except json.JSONDecodeError: + raise RuntimeError( + f"Failed to decode JSON from config file in {model}. Please ensure the file is valid JSON." + ) + + if raw_config.get("model_type", "llama") in _CONFIG_REGISTRY: + config_class = _CONFIG_REGISTRY[raw_config["model_type"]] + config = config_class.from_pretrained(model, revision=revision) + setattr(config, "_name_or_path", model) + else: + try: + config = AutoConfig.from_pretrained( + model, trust_remote_code=trust_remote_code, revision=revision, **kwargs + ) + except ValueError as e: + raise e + + _materialize_architectures(config, raw_config) + _restore_raw_glm_dsa_fields(config, raw_config) + + # extract 'text_config' + text_config = get_hf_text_config(config) + + # quantization config will copy to text_config + if hasattr(text_config, "quantization_config"): + if "modules_to_not_convert" in text_config.quantization_config: + text_config.quantization_config["ignored_layers"] = ( + text_config.quantization_config["modules_to_not_convert"] + ) + del text_config.quantization_config["modules_to_not_convert"] + + # If the draft head ships in the same checkpoint as the base model, + # rewrite the architecture in place so the model loader dispatches + # to the *NextN / *Eagle3 entry class instead of the base one. + # ``architectures`` is guaranteed non-None here when the on-disk + # config.json declared it (see the source-of-truth pin above); + # the truthiness check stays for configs that genuinely lack the + # field. + if ( + is_draft_worker + and config.architectures + and "NextN" not in config.architectures[0] + and "Eagle" not in config.architectures[0] + and "DFlash" not in config.architectures[0] + ): + if config.architectures[0] == "MiniMaxM2ForCausalLM": + config.architectures[0] = "LlamaForCausalLMEagle3" + else: + config.architectures[0] += "NextN" + + if text_config.architectures == ["LlamaForCausalLMNextN"]: + text_config.num_hidden_layers = 1 + + if model_override_args: + text_config.update(model_override_args) + + if resolve_architecture(config) in [ + "KimiK25ForConditionalGeneration", + "KimiK25Config", + "Qwen3_5MoeForConditionalGeneration", + "Qwen3_5MoeForConditionalGenerationNextN", + "Qwen3_5MoeConfig", + "Qwen3_5ForConditionalGeneration", + "Qwen3_5ForConditionalGenerationNextN", + "Qwen3OmniMoeForConditionalGeneration", + "Qwen3OmniMoeConfig", + "Qwen3ASRForConditionalGeneration", + "Qwen3ASRConfig", + ]: + config.text_config = text_config + return config + + return text_config + + +@lru_cache_frozenset(maxsize=32) +def get_generation_config( + model: str, + trust_remote_code: bool, + revision: str | None = None, + **kwargs, +): + try: + return GenerationConfig.from_pretrained( + model, trust_remote_code=trust_remote_code, revision=revision, **kwargs + ) + except OSError: + logging.debug("model doesn't have generation_config.json") + return None + + +# Models don't use the same configuration key for determining the maximum +# context length. Store them here so we can sanely check them. +# 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" + + +# Architectures for which ``tokenizer.json`` encodes the exact pre-tokenizer +# / normalizer the model was trained with, and whose AutoTokenizer defaults +# diverge from that. Kimi-K2.5 ships a custom ``TikTokenTokenizer`` via +# ``trust_remote_code`` that AutoTokenizer already handles correctly, so this +# verbatim tokenizer path must stay architecture-gated. +_VERBATIM_TOKENIZER_ARCHITECTURES: frozenset = frozenset( + { + "MiniMaxM2ForCausalLM", + } +) +_DEEPSEEK_V4_TOKENIZER_ARCHITECTURES: frozenset = frozenset( + { + "DeepseekV4ForCausalLM", + } +) + + +def prefers_verbatim_fast_tokenizer(architectures: list[str] | None) -> bool: + """True if the model's architectures warrant bypassing AutoTokenizer and + loading ``PreTrainedTokenizerFast`` from ``tokenizer.json`` verbatim. + """ + if not architectures: + return False + return any(arch in _VERBATIM_TOKENIZER_ARCHITECTURES for arch in architectures) + + +def prefers_deepseek_v4_tokenizer(architectures: list[str] | None) -> bool: + if not architectures: + return False + return any(arch in _DEEPSEEK_V4_TOKENIZER_ARCHITECTURES for arch in architectures) + + +def _find_deepseek_v4_encoding_file( + tokenizer_name: str, + tokenizer_revision: str | None, +) -> str: + if os.path.isdir(tokenizer_name): + encoding_path = os.path.join(tokenizer_name, "encoding", "encoding_dsv4.py") + if os.path.exists(encoding_path): + return encoding_path + + try: + encoding_path = cached_file( + tokenizer_name, + "encoding/encoding_dsv4.py", + revision=tokenizer_revision, + _raise_exceptions_for_gated_repo=False, + _raise_exceptions_for_missing_entries=False, + _raise_exceptions_for_connection_errors=False, + ) + except TypeError: + encoding_path = cached_file( + tokenizer_name, + "encoding/encoding_dsv4.py", + revision=tokenizer_revision, + ) + + if not encoding_path: + raise RuntimeError( + "DeepSeek V4 tokenizer mode requires " + "`encoding/encoding_dsv4.py` from the model repository." + ) + return encoding_path + + +def _load_deepseek_v4_encode_messages( + tokenizer_name: str, + tokenizer_revision: str | None, +) -> Callable[..., str]: + encoding_path = _find_deepseek_v4_encoding_file(tokenizer_name, tokenizer_revision) + spec = importlib.util.spec_from_file_location( + _DEEPSEEK_V4_ENCODING_MODULE_NAME, encoding_path + ) + if spec is None or spec.loader is None: + raise RuntimeError(f"Unable to load DeepSeek V4 encoding from {encoding_path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + encode_messages = getattr(module, "encode_messages", None) + if encode_messages is None: + raise RuntimeError(f"{encoding_path} does not define encode_messages") + return encode_messages + + +def _wrap_deepseek_v4_tokenizer( + tokenizer: PreTrainedTokenizer | PreTrainedTokenizerFast, + encode_messages: Callable[..., str], +) -> PreTrainedTokenizer | PreTrainedTokenizerFast: + """Attach DeepSeek V4's model-provided chat encoder to a HF tokenizer. + + This loads the official encoder from the checkpoint instead of vendoring it + in TokenSpeed. + """ + + dsv4_tokenizer = copy.copy(tokenizer) + added_vocab = tokenizer.get_added_vocab() + added_vocab_size = len(added_vocab) + tokenizer_vocab_size = tokenizer.vocab_size + + class _DeepseekV4Tokenizer(tokenizer.__class__): # type: ignore + def apply_chat_template( + self, + messages: list[dict[str, Any]], + tools: list[dict[str, Any]] | None = None, + **kwargs, + ): + thinking = kwargs.get("thinking", False) or kwargs.get( + "enable_thinking", False + ) + conversation = kwargs.get("conversation", messages) + conversation = conversation.copy() + if tools: + conversation.insert(0, {"role": "system", "tools": tools}) + + reasoning_effort = kwargs.get("reasoning_effort") + if reasoning_effort not in ("max", "high"): + reasoning_effort = None + + prompt = encode_messages( + conversation, + thinking_mode="thinking" if thinking else "chat", + drop_thinking=kwargs.get("drop_thinking", True), + reasoning_effort=reasoning_effort, + ) + + if not kwargs.get("tokenize", True): + return prompt + + return_dict = kwargs.get("return_dict", False) + forwarded_keys = ( + "truncation", + "max_length", + "padding", + "return_tensors", + "return_attention_mask", + "return_token_type_ids", + "return_special_tokens_mask", + "return_offsets_mapping", + "return_length", + ) + forwarded = {k: kwargs[k] for k in forwarded_keys if k in kwargs} + encoding = self(prompt, add_special_tokens=False, **forwarded) + if return_dict: + return encoding + return encoding["input_ids"] + + def num_special_tokens_to_add(self) -> int: + return len(self.encode("")) + + def __len__(self) -> int: + return tokenizer_vocab_size + added_vocab_size + + def get_added_vocab(self) -> dict[str, int]: + return added_vocab.copy() + + _DeepseekV4Tokenizer.__name__ = f"DSV4{tokenizer.__class__.__name__}" + dsv4_tokenizer.__class__ = _DeepseekV4Tokenizer + return dsv4_tokenizer + + +def get_tokenizer( + tokenizer_name: str, + *args, + tokenizer_mode: str = "auto", + trust_remote_code: bool = False, + tokenizer_revision: str | None = None, + architectures: list[str] | None = None, + **kwargs, +) -> PreTrainedTokenizer | PreTrainedTokenizerFast: + """Gets a tokenizer for the given model name via Huggingface. + + ``architectures`` is the model's ``config.architectures`` list (caller + should pass it when available). It gates whether we bypass AutoTokenizer + and load ``PreTrainedTokenizerFast`` from ``tokenizer.json`` verbatim — + needed for a small set of models (e.g. MiniMax-M2) whose AutoTokenizer + defaults diverge from training. Models with custom tokenizer classes + loaded via ``trust_remote_code`` (e.g. Kimi-K2.5's ``TikTokenTokenizer``) + must NOT go through the verbatim path; leaving ``architectures`` as None + (the default) keeps the safe AutoTokenizer-only behavior. + """ + 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 + + fast_tokenizer = None + if ( + tokenizer_mode != "slow" + and kwargs.get("use_fast", True) + and prefers_verbatim_fast_tokenizer(architectures) + ): + try: + fast_tokenizer = PreTrainedTokenizerFast.from_pretrained( + tokenizer_name, + *args, + revision=tokenizer_revision, + clean_up_tokenization_spaces=False, + ) + except Exception: + fast_tokenizer = None + + 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 + + # Swap in the fast tokenizer, carrying over chat_template from + # tokenizer_config.json if tokenizer.json doesn't have one. + if fast_tokenizer is not None and fast_tokenizer is not tokenizer: + if getattr(tokenizer, "chat_template", None) and not getattr( + fast_tokenizer, "chat_template", None + ): + fast_tokenizer.chat_template = tokenizer.chat_template + tokenizer = fast_tokenizer + + if not isinstance(tokenizer, PreTrainedTokenizerFast): + warnings.warn( + "Using a slow tokenizer. This might cause a significant " + "slowdown. Consider using a fast tokenizer instead." + ) + + if tokenizer_mode == "auto" and prefers_deepseek_v4_tokenizer(architectures): + tokenizer = _wrap_deepseek_v4_tokenizer( + tokenizer, + _load_deepseek_v4_encode_messages(tokenizer_name, tokenizer_revision), + ) + + attach_additional_stop_token_ids(tokenizer) + return tokenizer + + +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 diff --git a/python/tokenspeed/runtime/utils/hostfunc.py b/python/tokenspeed/runtime/utils/hostfunc.py new file mode 100644 index 0000000..82a0feb --- /dev/null +++ b/python/tokenspeed/runtime/utils/hostfunc.py @@ -0,0 +1,220 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Host-function CUDA graph node. + +Wraps ``cudaLaunchHostFunc`` so Python callables can be recorded as nodes +inside a CUDA graph. The C++ stub is JIT-compiled on first use via +``torch.utils.cpp_extension.load_inline``. +""" + +from __future__ import annotations + +import atexit +import threading +from collections.abc import Callable +from typing import Any + +import torch + +_ext = None +_ext_lock = threading.Lock() + + +_CPP_SRC = r""" +#include +#include +#include +#include +#include +#include + +namespace py = pybind11; + +struct HostFuncUserData { + bool free_user_data; // trampoline deletes self when true + py::function fn; + py::tuple args; + py::dict kwargs; + + HostFuncUserData(bool fud, py::function f, py::tuple a, py::dict k) + : free_user_data(fud), fn(std::move(f)), + args(std::move(a)), kwargs(std::move(k)) {} +}; + +static void CUDART_CB host_func_trampoline(void* user_data) { + // cudaLaunchHostFunc runs on a driver thread without the GIL. + py::gil_scoped_acquire gil; + auto* data = static_cast(user_data); + try { + data->fn(*data->args, **data->kwargs); + } catch (const py::error_already_set& e) { + // Can't throw across the CUDA callback boundary; log + swallow. + PyErr_Print(); + } catch (const std::exception& e) { + PyErr_SetString(PyExc_RuntimeError, e.what()); + PyErr_Print(); + } + if (data->free_user_data) { + delete data; + } + // else: caller keeps the handle and frees it explicitly. +} + +// stream_ptr: CUstream as uintptr_t (torch.cuda.Stream.cuda_stream) +// free_user_data: True for non-capturing stream (safe to free after call); +// False during capture (callback replays from the same +// HostFuncUserData, so caller must keep it alive). +// Returns the user-data pointer as int when free_user_data=False, else None. +// Called with the GIL held (pybind11 default). We only release GIL for +// the actual cudaLaunchHostFunc call. +static std::optional launch_hostfunc( + uintptr_t stream_ptr, + bool free_user_data, + py::function fn, + py::args args, + py::kwargs kwargs) +{ + auto* data = new HostFuncUserData( + free_user_data, std::move(fn), py::tuple(args), py::dict(kwargs)); + cudaStream_t stream = reinterpret_cast(stream_ptr); + cudaError_t err; + { + py::gil_scoped_release release; + err = cudaLaunchHostFunc(stream, host_func_trampoline, data); + } + if (err != cudaSuccess) { + delete data; + throw std::runtime_error( + std::string("cudaLaunchHostFunc failed: ") + cudaGetErrorString(err)); + } + if (free_user_data) { + return std::nullopt; + } + return reinterpret_cast(data); +} + +static void free_hostfunc_user_data(uintptr_t handle) { + auto* data = reinterpret_cast(handle); + delete data; +} + +PYBIND11_MODULE(tokenspeed_hostfunc_ext, m) { + m.def("launch_hostfunc", &launch_hostfunc, + "Schedule a Python callable on a CUDA stream via cudaLaunchHostFunc."); + m.def("free_hostfunc_user_data", &free_hostfunc_user_data, + "Free user data reserved for a captured host function."); +} +""" + + +def _load_ext(): + global _ext + if _ext is not None: + return _ext + with _ext_lock: + if _ext is not None: + return _ext + import os + + from torch.utils.cpp_extension import CUDA_HOME, load_inline + + # Resolve CUDA headers explicitly — PyTorch's CUDA_HOME discovery + # sometimes misses conda-forge's ``$PREFIX/targets/.../include`` layout. + cuda_home = CUDA_HOME or os.environ.get("CUDA_HOME") + include_dirs = [] + library_dirs = [] + if cuda_home: + inc = os.path.join(cuda_home, "include") + lib = os.path.join(cuda_home, "lib64") + if os.path.isdir(inc): + include_dirs.append(inc) + if os.path.isdir(lib): + library_dirs.append(lib) + for cand in ( + "/home/jue/miniforge3/targets/x86_64-linux/include", + "/opt/conda/targets/x86_64-linux/include", + ): + if os.path.isdir(cand) and cand not in include_dirs: + include_dirs.append(cand) + + _ext = load_inline( + name="tokenspeed_hostfunc_ext", + cpp_sources=[_CPP_SRC], + extra_cflags=["-O2"] + [f"-I{d}" for d in include_dirs], + extra_ldflags=[f"-L{d}" for d in library_dirs] + ["-lcudart"], + verbose=False, + ) + return _ext + + +# Handles kept alive for the lifetime of a captured graph — freeing them +# before the graph is destroyed would dangle-ref inside the trampoline. +_USER_DATA_HANDLES: set[int] = set() + + +def launch_hostfunc(fn: Callable, *args: Any, **kwargs: Any) -> int | None: + """Run ``fn(*args, **kwargs)`` on the current CUDA stream. + + When capturing, returns a handle to the user-data the caller must keep + alive; otherwise executes eagerly and returns None. + """ + stream = torch.cuda.current_stream() + is_capturing = torch.cuda.is_current_stream_capturing() + ext = _load_ext() + handle = ext.launch_hostfunc( + stream.cuda_stream, not is_capturing, fn, *args, **kwargs + ) + if is_capturing: + if handle is None: + raise RuntimeError("hostfunc launch did not return a capture handle.") + _USER_DATA_HANDLES.add(handle) + else: + if handle is not None: + raise RuntimeError("hostfunc eager launch unexpectedly returned a handle.") + return handle + + +def hostfunc(fn: Callable) -> Callable: + """Decorator: turn ``fn`` into a host-function graph node when called.""" + + def wrapper(*args: Any, **kwargs: Any): + return launch_hostfunc(fn, *args, **kwargs) + + wrapper.__wrapped__ = fn # type: ignore[attr-defined] + return wrapper + + +def free_hostfunc_user_data(handle: int) -> None: + if handle not in _USER_DATA_HANDLES: + raise ValueError(f"hostfunc user-data handle {handle} not tracked") + _load_ext().free_hostfunc_user_data(handle) + _USER_DATA_HANDLES.remove(handle) + + +def free_all_hostfunc_user_data() -> None: + if _ext is None: + return + for handle in list(_USER_DATA_HANDLES): + _ext.free_hostfunc_user_data(handle) + _USER_DATA_HANDLES.clear() + + +atexit.register(free_all_hostfunc_user_data) diff --git a/python/tokenspeed/runtime/utils/network.py b/python/tokenspeed/runtime/utils/network.py new file mode 100644 index 0000000..29cd479 --- /dev/null +++ b/python/tokenspeed/runtime/utils/network.py @@ -0,0 +1,130 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Port, socket, and host networking helpers.""" + +from __future__ import annotations + +import logging +import os +import socket +import warnings + +logger = logging.getLogger(__name__) + + +def is_port_available(port): + """Return whether a port is available.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + try: + s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + s.bind(("", port)) + s.listen(1) + return True + except OSError: + return False + except OverflowError: + return False + + +def get_free_port(): + try: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + return s.getsockname()[1] + except OSError: + with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + return s.getsockname()[1] + + +def set_uvicorn_logging_configs(): + from uvicorn.config import LOGGING_CONFIG + + LOGGING_CONFIG["formatters"]["default"][ + "fmt" + ] = "[%(asctime)s] %(levelprefix)s %(message)s" + LOGGING_CONFIG["formatters"]["default"]["datefmt"] = "%Y-%m-%d %H:%M:%S" + LOGGING_CONFIG["formatters"]["access"][ + "fmt" + ] = '[%(asctime)s] %(levelprefix)s %(client_addr)s - "%(request_line)s" %(status_code)s' + LOGGING_CONFIG["formatters"]["access"]["datefmt"] = "%Y-%m-%d %H:%M:%S" + + +def get_local_ip_by_remote() -> str | None: + s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + try: + s.connect(("8.8.8.8", 80)) + return s.getsockname()[0] + except Exception: + pass + + try: + hostname = socket.gethostname() + ip = socket.gethostbyname(hostname) + if ip and ip != "127.0.0.1" and ip != "0.0.0.0": + return ip + except Exception: + pass + + try: + s = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM) + s.connect(("2001:4860:4860::8888", 80)) + return s.getsockname()[0] + except Exception: + raise ValueError("Can not get local ip") + + +def get_ip() -> str: + from tokenspeed.runtime.utils.env import envs + + host_ip = envs.TOKENSPEED_HOST_IP.get() or os.getenv("HOST_IP", "") + if host_ip: + return host_ip + + try: + hostname = socket.gethostname() + ip = socket.gethostbyname(hostname) + if ip and ip != "127.0.0.1" and ip != "0.0.0.0": + return ip + except Exception: + pass + + s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + try: + s.connect(("8.8.8.8", 80)) + return s.getsockname()[0] + except Exception: + pass + + try: + s = socket.socket(socket.AF_INET6, socket.SOCK_DGRAM) + s.connect(("2001:4860:4860::8888", 80)) + return s.getsockname()[0] + except Exception: + pass + + warnings.warn( + "Failed to get the IP address, using 0.0.0.0 by default." + "The value can be set by the environment variable" + " TOKENSPEED_HOST_IP or HOST_IP.", + stacklevel=2, + ) + return "0.0.0.0" diff --git a/python/tokenspeed/runtime/utils/nvtx.py b/python/tokenspeed/runtime/utils/nvtx.py new file mode 100644 index 0000000..fd51f92 --- /dev/null +++ b/python/tokenspeed/runtime/utils/nvtx.py @@ -0,0 +1,121 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Zero-overhead NVTX range wrapper. + +Off by default. Enable via `--enable-nvtx` or `TOKENSPEED_NVTX=1`. +`nvtx_range(...)` is dual-purpose — context manager or decorator: + + @nvtx_range("forward", color="blue") + def forward(self, ...): ... + + with nvtx_range("sampling", color="yellow"): + backend.sample(...) + +The `_enabled` check is re-read at each call, so a decorator applied at +import time still arms once the CLI flag flips the flag at runtime. +""" + +from __future__ import annotations + +import functools + +from tokenspeed_kernel.platform import current_platform + +from tokenspeed.runtime.utils.env import envs + +if current_platform().is_nvidia: + import nvtx + from nvtx.colors import _NVTX_COLORS +else: + nvtx = None + _NVTX_COLORS = () + +NVTX_DOMAIN: str = "tokenspeed" + +# Built-in colors — keeping tokenspeed self-contained (no matplotlib +# dependency for color resolution). Anything else falls back to "blue" +# rather than crashing the server. +_VALID_COLORS = frozenset(c for c in _NVTX_COLORS if c is not None) + +_nvtx_available = nvtx is not None +_enabled: bool = _nvtx_available and envs.TOKENSPEED_NVTX.get() + + +def set_nvtx_enabled(enabled: bool) -> None: + global _enabled + _enabled = _nvtx_available and enabled + + +class _Range: + __slots__ = ("name", "color", "category") + + def __init__( + self, + name: str, + color: str, + category: str | None, + ) -> None: + self.name = name + self.color = color + self.category = category + + def __enter__(self): + if _enabled: + nvtx.push_range( + message=self.name, + color=self.color, + domain=NVTX_DOMAIN, + category=self.category, + ) + return self + + def __exit__(self, *exc): + if _enabled: + nvtx.pop_range(domain=NVTX_DOMAIN) + return False + + def __call__(self, func): + name, color, category = self.name, self.color, self.category + + @functools.wraps(func) + def wrapper(*args, **kwargs): + if not _enabled: + return func(*args, **kwargs) + nvtx.push_range( + message=name, color=color, domain=NVTX_DOMAIN, category=category + ) + try: + return func(*args, **kwargs) + finally: + nvtx.pop_range(domain=NVTX_DOMAIN) + + return wrapper + + +def nvtx_range( + name: str, + *, + color: str = "blue", + category: str | None = None, +) -> _Range: + if color not in _VALID_COLORS: + color = "blue" + return _Range(name, color, category) diff --git a/python/tokenspeed/runtime/utils/pdl.py b/python/tokenspeed/runtime/utils/pdl.py new file mode 100644 index 0000000..d772382 --- /dev/null +++ b/python/tokenspeed/runtime/utils/pdl.py @@ -0,0 +1,36 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from functools import lru_cache + +from tokenspeed_kernel.platform import current_platform + + +@lru_cache(maxsize=1) +def pdl_enabled() -> bool: + """Return whether Programmatic Dependent Launch is enabled.""" + from tokenspeed.runtime.utils.env import global_server_args_dict + + if global_server_args_dict.get("disable_pdl", False): + return False + try: + return current_platform().is_hopper_plus + except Exception: + return False diff --git a/python/tokenspeed/runtime/utils/process.py b/python/tokenspeed/runtime/utils/process.py new file mode 100644 index 0000000..bd8d0f7 --- /dev/null +++ b/python/tokenspeed/runtime/utils/process.py @@ -0,0 +1,78 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""Process and signal helpers.""" + +from __future__ import annotations + +import logging +import os +import signal +import sys +import threading + +import psutil + +logger = logging.getLogger(__name__) + + +def register_usr_signal(): + parent_process = psutil.Process().parent() + + def signal_handler(sig, frame): + logger.error("recv usr signal, kill usr signal to parent") + parent_process.send_signal(signal.SIGUSR1) + + signal.signal(signal.SIGUSR1, signal_handler) + + +def kill_process_tree(parent_pid, include_parent: bool = True, skip_pid: int = None): + """Kill the target process and all of its child processes.""" + if threading.current_thread() is threading.main_thread(): + signal.signal(signal.SIGCHLD, signal.SIG_DFL) + + if parent_pid is None: + parent_pid = os.getpid() + include_parent = False + + try: + itself = psutil.Process(parent_pid) + except psutil.NoSuchProcess: + return + + children = itself.children(recursive=True) + for child in children: + if child.pid == skip_pid: + continue + try: + child.kill() + except psutil.NoSuchProcess: + pass + + if include_parent: + try: + if parent_pid == os.getpid(): + itself.kill() + sys.exit(0) + + itself.kill() + itself.send_signal(signal.SIGQUIT) + except psutil.NoSuchProcess: + pass diff --git a/python/tokenspeed/runtime/utils/server_args.py b/python/tokenspeed/runtime/utils/server_args.py new file mode 100755 index 0000000..7394b8e --- /dev/null +++ b/python/tokenspeed/runtime/utils/server_args.py @@ -0,0 +1,1990 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +"""The arguments of the server.""" + +import argparse +import dataclasses +import json +import os +import random +from typing import Literal + +from tokenspeed_kernel.ops.attention.triton.linear.chunk_delta_h import ( + CHUNK_SIZE as FLA_CHUNK_SIZE, +) +from tokenspeed_kernel.platform import current_platform + +from tokenspeed.runtime.distributed.mapping import Mapping, _resolve_parallelism_sizes +from tokenspeed.runtime.utils import ( + get_amdgpu_memory_capacity, + get_colorful_logger, + get_nvgpu_memory_capacity, + is_valid_ipv6_address, + maybe_model_redirect, + nullable_str, +) +from tokenspeed.runtime.utils.network import is_port_available + +logger = get_colorful_logger(__name__) + +ENABLE_CP = os.environ.get("ENABLE_CP", "false").lower() in ("true", "1") + + +def str_to_bool(value: str | bool) -> bool: + if isinstance(value, bool): + return value + normalized = value.lower() + if normalized in {"1", "true", "yes", "on"}: + return True + if normalized in {"0", "false", "no", "off"}: + return False + raise argparse.ArgumentTypeError(f"invalid boolean value: {value!r}") + + +@dataclasses.dataclass +class ServerArgs: + # Model and tokenizer + model: str + tokenizer: str | None = None + tokenizer_mode: str = "auto" + skip_tokenizer_init: bool = False + load_format: str = "auto" + trust_remote_code: bool = True + dtype: str = "auto" + kv_cache_dtype: str = "auto" + kv_cache_quant_method: str = "none" + quantization: str | None = None + quantization_param_path: nullable_str = None + max_model_len: int | None = None + device: str = "cuda" + served_model_name: str | None = None + revision: str | None = None + language_model_only: bool = False + + # Port for the HTTP server + host: str = "127.0.0.1" + port: int = 8000 + + # Memory and scheduling + gpu_memory_utilization: float | None = None + max_num_seqs: int | None = None + max_total_tokens: int | None = None + chunked_prefill_size: int | None = None + max_prefill_tokens: int = 8192 + enable_mixed_batch: bool = False + block_size: int = 64 + # special kv cache + mamba_ssm_dtype: str = "float32" + mamba_track_interval: int = 256 + max_mamba_cache_size: int | None = None + mamba_full_memory_ratio: float = 0.9 + enable_mamba_l2: bool = False + mamba_l2_host_slots: int = 0 + mamba_l2_ratio: float = 2.0 + mamba_l2_layout: str = "layer_first" + mamba_l2_io_backend: str = "kernel" + mamba_l2_host_gb: int = 0 + + # Other runtime options + stream_interval: int = 1 + stream_output: bool = False + # Inline detokenization is the only supported path and is intentionally + # not configurable from the CLI. + enable_inline_detokenizer: bool = True + seed: int | None = None + distributed_timeout_seconds: int | None = None + download_dir: str | None = None + # Used for customizing extensible models + ext_yaml: str | None = None + base_gpu_id: int = 0 + gpu_id_step: int = 1 + + # Logging + log_level: str = "info" + log_level_http: str | None = None + enable_log_requests: bool = False + log_requests_level: int = 0 + enable_log_request_stats: bool = False + enable_metrics: bool = False + decode_log_interval: int = 40 + metrics_reporters: list[str] | None = None + app_key: str | None = None + + # API related + api_key: str | None = None + enable_cache_report: bool = False + kv_events_config: str | None = None + + # Data parallelism + data_parallel_size: int | None = None + load_balance_method: str = "shortest_queue" + load_watch_interval: float = 0.02 + + # Expert parallelism + ep_size: int = 1 + init_expert_location: str = "trivial" + ep_num_redundant_experts: int = 0 + ep_dispatch_algorithm: ( + Literal[ + "static", + "dynamic", + "fake", + "static_with_zero_expert", + "dynamic_with_zero_expert", + ] + | None + ) = None + eplb_algorithm: str = "auto" + expert_distribution_recorder_mode: ( + Literal["stat", "stat_approx", "per_pass", "per_token"] | None + ) = None + expert_distribution_recorder_buffer_size: int | None = None + enable_expert_distribution_metrics: bool = False + enable_eplb: bool = False + + # MoE backend + moe_backend: str = "auto" + draft_moe_backend: str | None = None + all2all_backend: str = "none" + deepep_mode: Literal["auto", "normal", "low_latency"] = "auto" + disable_flashinfer_cutlass_moe_fp4_allgather: bool = False + + # KVStore + enable_kvstore: bool = False + kvstore_ratio: float = 2.0 + kvstore_size: int = 0 + kvstore_io_backend: str = "kernel" + kvstore_mem_layout: str = "layer_first" + kvstore_storage_backend: str | None = None + kvstore_storage_backend_extra_config: str | None = None + enable_mla_l1_5_cache: bool = False + + # Multi-node distributed serving + dist_init_addr: str | None = None + nnodes: int = 1 + node_rank: int = 0 + + # Hugging Face model config overrides in JSON + hf_overrides: str = "{}" + preferred_sampling_params: str | None = None + + # Kernel backend + attention_backend: str | None = None + drafter_attention_backend: str | None = None + sampling_backend: str | None = None + dp_sampling: bool = False + dp_sampling_min_bs: int | None = None + attention_use_fp4_indexer_cache: bool | None = None + use_trtllm_ragged_deepseek_prefill: bool | None = None + + # DeepSeek V4 + deepseek_v4_mega_moe_max_num_tokens: int = 0 + deepseek_v4_indexer_prefill_max_logits_mb: int = 512 + deepseek_v4_prefill_chunk_size: int = 4 + + # Grammar backend + grammar_backend: str = "none" + # Used by ``input_processor`` to defer json_schema grammars past the + # model's reasoning channel. + reasoning_parser: str | None = None + grammar_compile_timeout_secs: float = 30.0 + grammar_compile_max_retries: int = 2 + disable_any_whitespace: bool = False + # Force the synchronous eager grammar fallback even on CUDA. Useful + # for parity-testing against the captured-grammar path (output should + # match; throughput will be lower since the sync stalls every step). + disable_capturable_grammar: bool = False + + # Speculative decoding + draft_model_path_use_base: bool | None = False + speculative_config: str | None = None + speculative_algorithm: str | None = None + speculative_draft_model_path: str | None = None + speculative_draft_model_quantization: str | None = "unquant" + speculative_num_steps: int = 3 + speculative_eagle_topk: int = 1 + speculative_num_draft_tokens: int | None = None + eagle3_layers_to_capture: str | None = None + # Logprob support flags — all OFF by default. Enabling extends the + # captured CUDA-graph footprint; requests asking for logprobs on a + # server started without the matching flag will receive empty logprobs. + enable_output_logprobs: bool = False + + # Runtime options + disable_pdl: bool = False + enable_prefix_caching: bool = True + disable_kvstore: bool = False + enforce_eager: bool = False + disable_cuda_graph_padding: bool = False + enable_cudagraph_gc: bool = False + enable_nccl_nvls: bool = False + enable_symm_mem: bool = False + disable_custom_all_reduce: bool = False + disable_overlap_schedule: bool = False + disable_tf32: bool = False + force_deterministic_rsag: bool = False + disable_sampling_tp_sync: bool = False + low_latency_max_num_tokens_per_gpu: int = 256 + max_cudagraph_capture_size: int | None = None + disable_prefill_graph: bool | None = False + # Breakable prefill graph bucket cap: None = auto min(2048, chunk); 0 disables. + prefill_graph_max_tokens: int | None = None + # Explicit prefill bucket list; unset = the relative-stride ladder (see get_prefill_token_buckets). + prefill_graph_capture_sizes: list[int] | None = None + cudagraph_capture_sizes: list[int] | None = None + enable_nan_detection: bool = False + enable_nvtx: bool = False + enable_p2p_check: bool = False + triton_attention_reduce_in_fp32: bool = False + delete_ckpt_after_loading: bool = False + weight_loader_prefetch_checkpoints: bool = False + weight_loader_prefetch_num_threads: int = 4 + enable_memory_saver: bool = False + enable_custom_logit_processor: bool = False + mla_disable_ragged: bool = False + warmups: str | None = None + + # parallel strategy + nprocs_per_node: int | None = None + world_size: int | None = None + attn_tp_size: int | None = None + dense_tp_size: int | None = None + moe_tp_size: int | None = None + mapping: Mapping | None = None + + mla_chunk_multiplier: int = 4 + mm_attention_backend: str | None = None + + # For PD/EPD disaggregation: "null", "prefill", "decode", or "encode" (vision-tower-only). + disaggregation_mode: str = "null" + disaggregation_bootstrap_port: int = 8998 + disaggregation_transfer_backend: str = "mooncake" + disaggregation_ib_device: str | None = None + disaggregation_layerwise_interval: int = 1 + pdlb_url: str | None = None + + skip_server_warmup: bool = False + + # For communication + norm fusion + comm_fusion_max_num_tokens: int = 2048 + enable_allreduce_fusion: bool = False + + enable_expert_parallel: bool = False + + @property + def mamba_cache_chunk_size(self) -> int: + return max(FLA_CHUNK_SIZE, self.block_size) + + def __post_init__(self): + self.resolve_basic_defaults() + self.resolve_parallelism() + self.resolve_memory_and_scheduling() + self.resolve_kernel_backends() + self.resolve_cache() + self.resolve_speculative_decoding() + self.resolve_communication() + self.resolve_disaggregation() + self.validate() + + def resolve_basic_defaults(self): + self.model = maybe_model_redirect(self.model) + + if self.kv_cache_dtype == "fp8": + self.kv_cache_dtype = "fp8_e4m3" + + self.resolve_config_aliases() + + # Set missing default values + if self.tokenizer is None: + self.tokenizer = self.model + + if self.served_model_name is None: + self.served_model_name = self.model + + if self.seed is None: + self.seed = random.randint(0, 1 << 30) + + def resolve_config_aliases(self): + if self.use_trtllm_ragged_deepseek_prefill is not None: + self.mla_disable_ragged = not self.use_trtllm_ragged_deepseek_prefill + + if self.speculative_config is not None: + try: + config = json.loads(self.speculative_config) + except json.JSONDecodeError as exc: + raise ValueError("--speculative-config must be valid JSON") from exc + + if not isinstance(config, dict): + raise ValueError("--speculative-config must be a JSON object") + + method = config.get("method") + if method is not None and self.speculative_algorithm is None: + self.speculative_algorithm = str(method).upper() + + draft_model = config.get("model") + if draft_model is not None and self.speculative_draft_model_path is None: + self.speculative_draft_model_path = str(draft_model) + + num_speculative_tokens = config.get("num_speculative_tokens") + if num_speculative_tokens is not None: + num_speculative_tokens = int(num_speculative_tokens) + if self.speculative_algorithm == "DFLASH": + if self.speculative_num_draft_tokens is None: + self.speculative_num_draft_tokens = num_speculative_tokens + self.speculative_num_steps = max(num_speculative_tokens - 1, 0) + else: + self.speculative_num_steps = num_speculative_tokens + + if self.speculative_num_draft_tokens is None: + self.speculative_num_draft_tokens = self.speculative_num_steps + 1 + + def resolve_memory_and_scheduling(self): + if current_platform().is_amd: + gpu_mem = get_amdgpu_memory_capacity() + elif current_platform().is_nvidia: + gpu_mem = get_nvgpu_memory_capacity() + else: + # GPU memory is not known yet or no GPU is available. + gpu_mem = None + + # Set GPU memory utilization, which depends on the tensor parallelism size. + self._gpu_memory_utilization_defaulted = False + if self.gpu_memory_utilization is None: + if self.mapping.world_size >= 16: + self.gpu_memory_utilization = 0.79 + elif self.mapping.world_size >= 8: + self.gpu_memory_utilization = 0.81 + elif self.mapping.world_size >= 4: + self.gpu_memory_utilization = 0.95 + elif self.mapping.world_size >= 2: + self.gpu_memory_utilization = 0.87 + else: + self.gpu_memory_utilization = 0.88 + self._gpu_memory_utilization_defaulted = True + + # Set the chunked prefill token budget. + if self.chunked_prefill_size is None: + self.chunked_prefill_size = 8192 + + # Set CUDA graph max capture size. + if self.max_cudagraph_capture_size is None: + # Based on detailed statistics, when serving TP1/TP2 models on lower-end GPUs with HBM<25G, you can either disable CUDA graph or set max_cudagraph_capture_size to a very small value to reduce graph memory overhead, with almost no impact on performance. TP4/TP8 serving still needs CUDA graph for high performance, and 80 is enough for lower-end GPUs. + if gpu_mem is not None and gpu_mem < 25_000: + if self.mapping.world_size < 4: + self.max_cudagraph_capture_size = 8 + else: + self.max_cudagraph_capture_size = 80 + elif self.speculative_algorithm: + self.max_cudagraph_capture_size = 80 + else: + self.max_cudagraph_capture_size = 160 + + # Set max number of sequences. + if self.max_num_seqs is None: + if self.speculative_algorithm: + self.max_num_seqs = 80 + else: + self.max_num_seqs = 160 + + def resolve_kernel_backends(self): + # Choose kernel backends + # attention_backend default is NOT set here — deferred to + # AttnInitializer.modify_args where both hardware and model arch are known. + + if self.sampling_backend is None: + # ``flashinfer`` is the only built-in backend that respects per-request + # ``temperature`` / ``top_p`` / ``top_k``. ``greedy`` is argmax-only + # (see ``GreedySamplingBackend.sample``: *"sampling_info is ignored + # for single-step (always argmax)"*) — fast for hand-tuned greedy + # decoding but silently wrong for any serving deployment where + # requests carry sampling params, since the model collapses into + # repetition-mode loops within a few hundred steps. Default to the + # sampling-respecting backend on NVIDIA where flashinfer is + # available, fall back to greedy elsewhere; users can still opt + # into greedy explicitly via ``--sampling-backend greedy``. + if current_platform().is_nvidia: + self.sampling_backend = "flashinfer" + else: + self.sampling_backend = "greedy" + + def resolve_parallelism(self): + world_size = self.world_size + nprocs_per_node = self.nprocs_per_node + nnodes = 1 if self.nnodes is None else self.nnodes + + attn_tp_size = self.attn_tp_size + attn_dp_size = self.data_parallel_size + + # ``ENABLE_CP`` interprets attention TP size as CP size. + attn_cp_size = 1 + if ENABLE_CP: + attn_cp_size, attn_tp_size = attn_tp_size, 1 + + if world_size is None: + world_size = 1 + if attn_tp_size is not None: + world_size *= attn_tp_size + if attn_cp_size is not None: + world_size *= attn_cp_size + if attn_dp_size is not None: + world_size *= attn_dp_size + logger.info( + "Inferred world_size (%s) from attn_tp_size (%s) x attn_cp_size (%s) x attn_dp_size (%s)", + world_size, + attn_tp_size, + attn_cp_size, + attn_dp_size, + ) + else: + logger.info("Specified world_size (%s)", world_size) + + attn_tp_size, attn_cp_size, attn_dp_size = _resolve_parallelism_sizes( + world_size, attn_tp_size, attn_cp_size, attn_dp_size + ) + + # Dense layers still default to full TP participation when no + # dedicated dense_tp_size is provided. + dense_tp_size = self.dense_tp_size + if self.dense_tp_size is None: + # dense always do tp now. + dense_tp_size = world_size + dense_dp_size = None + + # --enable-expert-parallel auto-sets ep_size = world_size + if self.enable_expert_parallel and self.ep_size == 1: + self.ep_size = world_size + logger.info("--enable-expert-parallel: auto-setting ep_size=%s", world_size) + + # MoE parallel sizes default to consuming the full world size unless + # the user overrides them explicitly. + moe_ep_size = 1 if self.ep_size is None else self.ep_size + moe_tp_size = ( + world_size // moe_ep_size if self.moe_tp_size is None else self.moe_tp_size + ) + moe_dp_size = None + + self.mapping = Mapping( + world_size=world_size, + attn_tp_size=attn_tp_size, + attn_cp_size=attn_cp_size, + attn_dp_size=attn_dp_size, + dense_tp_size=dense_tp_size, + dense_dp_size=dense_dp_size, + moe_tp_size=moe_tp_size, + moe_ep_size=moe_ep_size, + moe_dp_size=moe_dp_size, + nprocs_per_node=nprocs_per_node, + nnodes=nnodes, + base_gpu_id=self.base_gpu_id, + gpu_id_step=self.gpu_id_step, + ) + + # Impl constraints: + if self.mapping.moe.has_tp and self.mapping.moe.has_ep: + raise ValueError("MoE TP and EP cannot be both > 1") + + logger.info("Parallelism configuration:\n%s", self.mapping) + + def resolve_cache(self): + # Handle KVStore settings. + self._handle_kvstore() + self.validate_cache_options() + + def resolve_speculative_decoding(self): + # Keep drafter backend consistent with the main model unless explicitly set. + if ( + self.speculative_algorithm is not None + and self.drafter_attention_backend is None + ): + self.drafter_attention_backend = self.attention_backend + + if ( + self.speculative_algorithm == "MTP" + and self.speculative_draft_model_path is None + ): + self.draft_model_path_use_base = True + + if self.draft_model_path_use_base: + self.speculative_draft_model_path = self.model + + if self.speculative_draft_model_path == self.model: + self.draft_model_path_use_base = True + + if self.speculative_draft_model_quantization == "unquant": + self.speculative_draft_model_quantization = None + + if self.speculative_algorithm == "DFLASH": + expected_steps = max(int(self.speculative_num_draft_tokens) - 1, 0) + if self.speculative_num_steps == ServerArgs.speculative_num_steps: + self.speculative_num_steps = expected_steps + elif self.speculative_num_steps != expected_steps: + raise ValueError( + "DFLASH requires speculative_num_steps to equal " + "speculative_num_draft_tokens - 1. " + f"Got {self.speculative_num_steps=} and " + f"{self.speculative_num_draft_tokens=}." + ) + + if self.eagle3_layers_to_capture is not None: + self.eagle3_layers_to_capture = [ + int(x) for x in self.eagle3_layers_to_capture.split(",") + ] + + # Hoist the PD-decode topk == 1 check to startup. + if self.speculative_algorithm is not None and self.speculative_eagle_topk != 1: + raise ValueError( + "speculative_eagle_topk > 1 (tree spec) is not currently " + f"supported: {self.speculative_eagle_topk=}. Only chain spec " + "(topk=1) is wired end-to-end." + ) + + def resolve_communication(self): + # Auto-enable allreduce fusion on supported single-node TP configurations. + platform = current_platform() + if ( + not self.enable_allreduce_fusion + and (current_platform().is_hopper_plus or platform.is_amd) + and self.mapping.nnodes == 1 + and self.mapping.has_attn_tp + and not self.mapping.has_attn_dp + ): + self.enable_allreduce_fusion = True + logger.info("Auto-enabled allreduce fusion") + + if self.mapping.attn.tp_size != self.mapping.dense.tp_size: + self.comm_fusion_max_num_tokens = -1 + self.enable_allreduce_fusion = False + logger.info( + "allreduce is forbidden due to different attn_tp_size: %s and dense_tp_size: %s!", + self.mapping.attn.tp_size, + self.mapping.dense.tp_size, + ) + + def resolve_disaggregation(self): + # PD disaggregation + if self.disaggregation_mode == "prefill": + self.enforce_eager = True + logger.warning("CUDA graph is disabled for prefill server") + elif self.disaggregation_mode == "decode": + # Prefix caching stays configurable for decode servers. + logger.info( + "enable_prefix_caching=%r for decode server", + self.enable_prefix_caching, + ) + elif self.disaggregation_mode == "encode": + # Encode server: vision tower only, no LM / KV pool / prefix cache. + # enforce_eager left as-is (the vision tower keeps its own CUDA graph). + if self.mapping.has_attn_dp: + raise ValueError( + "disaggregation_mode=encode currently supports " + "data_parallel_size == 1 inside one encode server; run " + "multiple independent encode servers for horizontal scale." + ) + self.enable_prefix_caching = False + + # Prefill graph disable logic is handled by AttnInitializer.modify_args + # after the attention backend is resolved. + + if ( + self.disaggregation_mode == "prefill" + and self.load_balance_method != "round_robin" + ): + if self.mapping.has_attn_dp: + raise ValueError( + "Not supported when " + f"{self.disaggregation_mode=} {self.load_balance_method=} " + f"{self.mapping.attn.dp_size=}" + ) + + def _handle_kvstore(self): + if self.disaggregation_mode in ("decode", "encode"): + self.enable_kvstore = False + logger.info( + "%s instance has set enable_kvstore to False!", + self.disaggregation_mode, + ) + elif not self.disable_kvstore: + self.enable_kvstore = True + + if self.kvstore_storage_backend == "mooncake": + if self.kvstore_mem_layout == "layer_first": + self.kvstore_mem_layout = "page_first" + logger.warning( + "Mooncake storage backend does not support layer_first layout, switching to %s layout", + self.kvstore_mem_layout, + ) + + if self.kvstore_io_backend == "direct": + self.kvstore_io_backend = "kernel" + logger.warning( + "Mooncake storage backend uses page_first layout, which requires kernel io backend" + ) + + def validate_cache_options(self): + if self.enable_kvstore and not self.enable_prefix_caching: + raise ValueError( + "KVStore and disabled prefix caching are mutually exclusive " + "and cannot be used at the same time. Please use only one of them." + ) + + def validate(self): + if ( + self.max_num_seqs is not None + and self.max_num_seqs < self.mapping.attn.dp_size + ): + raise ValueError( + f"max_num_seqs must be >= attn_dp_size: {self.max_num_seqs=} < {self.mapping.attn.dp_size=}" + ) + + if self.mapping.has_attn_cp and self.max_num_seqs > 1: + raise ValueError("CP attention is enabled but max_num_seqs > 1") + + if self.mapping.has_attn_dp: + if self.chunked_prefill_size > self.max_prefill_tokens: + raise ValueError( + f"chunked_prefill_size must be <= max_prefill_tokens: {self.chunked_prefill_size=} > {self.max_prefill_tokens=}" + ) + + if self.deepseek_v4_prefill_chunk_size <= 0: + raise ValueError("deepseek_v4_prefill_chunk_size must be positive") + + if self.enable_eplb and (self.expert_distribution_recorder_mode is None): + self.expert_distribution_recorder_mode = "stat" + logger.info( + "EPLB is enabled. The expert_distribution_recorder_mode is automatically set." + ) + + if (self.enable_eplb or (self.init_expert_location is not None)) and ( + self.ep_dispatch_algorithm is None + ): + self.ep_dispatch_algorithm = "static" + logger.info( + "EPLB is enabled or init_expert_location is provided. ep_dispatch_algorithm is configured." + ) + + from tokenspeed.runtime.utils.env import envs + + envs.TOKENSPEED_MAMBA_SSM_DTYPE.set(self.mamba_ssm_dtype) + if not self.disable_pdl: + os.environ.setdefault("TORCHINDUCTOR_ENABLE_PDL", "1") + # Enable PDL for fused attention kernels. + os.environ.setdefault("TRTLLM_ENABLE_PDL", "1") + os.environ.setdefault("TLLM_LOG_LEVEL", "INFO") + + @staticmethod + def add_cli_args(parser: argparse.ArgumentParser): + parser.allow_abbrev = False + + # Model and port args + parser.add_argument( + "model_path", + nargs="?", + metavar="model", + default=None, + help="The model name or path (positional argument). " + "Equivalent to --model.", + ) + parser.add_argument( + "--model", + "--model-path", + metavar="MODEL", + type=str, + default=None, + help="The path of the model weights. This can be a local folder or a Hugging Face repo ID.", + ) + parser.add_argument( + "--tokenizer", + metavar="TOKENIZER", + type=str, + default=ServerArgs.tokenizer, + help="The path of the tokenizer.", + ) + parser.add_argument( + "--host", type=str, default=ServerArgs.host, help="The host of the server." + ) + parser.add_argument( + "--port", type=int, default=ServerArgs.port, help="The port of the server." + ) + parser.add_argument( + "--tokenizer-mode", + type=str, + default=ServerArgs.tokenizer_mode, + choices=["auto", "slow", "deepseek_v4"], + help="Tokenizer mode. 'auto' will use the fast " + "tokenizer and model-specific tokenizer hooks if available, " + "'slow' will always use the slow tokenizer.", + ) + parser.add_argument( + "--skip-tokenizer-init", + action=argparse.BooleanOptionalAction, + default=ServerArgs.skip_tokenizer_init, + help="If set, skip init tokenizer and pass input_ids in generate request", + ) + parser.add_argument( + "--language-model-only", + action="store_true", + default=ServerArgs.language_model_only, + help="Skip vision/audio encoders on a multimodal checkpoint and " + "run text-only. Multimodal requests are rejected.", + ) + parser.add_argument("--ext-yaml", type=str, default=None) + parser.add_argument( + "--load-format", + type=str, + default=ServerArgs.load_format, + choices=[ + "auto", + "pt", + "safetensors", + "npcache", + "dummy", + "extensible", + ], + help="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.', + ) + parser.add_argument( + "--trust-remote-code", + action=argparse.BooleanOptionalAction, + default=False, + help="Whether or not to allow for custom models defined on the Hub in their own modeling files.", + ) + parser.add_argument( + "--dtype", + type=str, + default=ServerArgs.dtype, + choices=["auto", "half", "float16", "bfloat16", "float", "float32"], + help="Data type for model weights and activations.\n\n" + '* "auto" will use FP16 precision for FP32 and FP16 models, and ' + "BF16 precision for BF16 models.\n" + '* "half" for FP16. Recommended for AWQ quantization.\n' + '* "float16" is the same as "half".\n' + '* "bfloat16" for a balance between precision and range.\n' + '* "float" is shorthand for FP32 precision.\n' + '* "float32" for FP32 precision.', + ) + parser.add_argument( + "--kv-cache-dtype", + type=str, + default=ServerArgs.kv_cache_dtype, + choices=["auto", "fp8", "fp8_e4m3"], + help='Data type for kv cache storage. "auto" will use model data type. "fp8" is an alias for "fp8_e4m3".', + ) + parser.add_argument( + "--kv-cache-quant-method", + type=str, + default=ServerArgs.kv_cache_quant_method, + choices=["none", "per_token_head"], + help="kv cache quant method", + ) + parser.add_argument( + "--quantization", + type=str, + default=ServerArgs.quantization, + choices=[ + "fp8", + "mxfp4", + "nvfp4", + "w8a8_fp8", + "compressed-tensors", + ], + help="The quantization method.", + ) + parser.add_argument( + "--quantization-param-path", + type=nullable_str, + default=None, + help="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. ", + ) + parser.add_argument( + "--max-model-len", + metavar="MAX_MODEL_LEN", + type=int, + default=ServerArgs.max_model_len, + help="The model's maximum context length. Defaults to None (will use the value from the model's config.json instead).", + ) + parser.add_argument( + "--device", + type=str, + default="cuda", + choices=["cuda"], + help="The device type.", + ) + parser.add_argument( + "--served-model-name", + type=str, + default=ServerArgs.served_model_name, + help="Override the model name returned by the v1/models endpoint in OpenAI API server.", + ) + parser.add_argument( + "--revision", + type=str, + default=None, + help="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.", + ) + # Memory and scheduling + parser.add_argument( + "--gpu-memory-utilization", + metavar="GPU_MEMORY_UTILIZATION", + type=float, + default=ServerArgs.gpu_memory_utilization, + help="The fraction of GPU memory to use for model weights and KV cache. Use a smaller value if you see out-of-memory errors.", + ) + parser.add_argument( + "--max-num-seqs", + metavar="MAX_NUM_SEQS", + type=int, + default=ServerArgs.max_num_seqs, + help="Maximum number of sequences to process concurrently.", + ) + parser.add_argument( + "--max-total-tokens", + type=int, + default=ServerArgs.max_total_tokens, + help="The maximum number of tokens in the memory pool. If not specified, it will be automatically calculated based on the memory usage fraction. " + "This overrides the automatically calculated token pool size.", + ) + parser.add_argument( + "--chunked-prefill-size", + metavar="CHUNKED_PREFILL_SIZE", + type=int, + default=ServerArgs.chunked_prefill_size, + help="Maximum number of tokens the scheduler may issue in a single iteration. Setting this to -1 disables chunked prefill.", + ) + parser.add_argument( + "--enable-mixed-batch", + action="store_true", + dest="enable_mixed_batch", + default=ServerArgs.enable_mixed_batch, + help="Allow the scheduler to issue prefill and decode requests in the same iteration.", + ) + parser.add_argument( + "--block-size", + metavar="BLOCK_SIZE", + type=int, + default=ServerArgs.block_size, + ) + + # KVStore + parser.add_argument( + "--disable-kvstore", + action="store_true", + help="Disable KVStore", + ) + parser.add_argument( + "--kvstore-ratio", + type=float, + default=ServerArgs.kvstore_ratio, + help="The ratio of the size of the KVStore host memory pool to the size of the device pool.", + ) + parser.add_argument( + "--kvstore-size", + type=int, + default=ServerArgs.kvstore_size, + help="The size of the KVStore host memory pool in gigabytes, which will override kvstore_ratio if set.", + ) + parser.add_argument( + "--kvstore-io-backend", + type=str, + choices=["direct", "kernel"], + default=ServerArgs.kvstore_io_backend, + help="The IO backend for KVStore transfer between CPU and GPU.", + ) + parser.add_argument( + "--kvstore-mem-layout", + type=str, + choices=[ + "layer_first", + "page_first", + "page_head", + ], + default=ServerArgs.kvstore_mem_layout, + help="The layout of the KVStore host memory pool.", + ) + parser.add_argument( + "--kvstore-storage-backend", + type=str, + choices=["mooncake"], + default=ServerArgs.kvstore_storage_backend, + help="The storage backend for KVStore. " + "Built-in backends: mooncake. " + "For dynamic backend, use --kvstore-storage-backend-extra-config to specify: " + "backend_name (custom name), module_path (Python module path), class_name (backend class name).", + ) + parser.add_argument( + "--kvstore-storage-backend-extra-config", + type=str, + default=ServerArgs.kvstore_storage_backend_extra_config, + help="A dictionary in JSON string format containing extra configuration for the storage backend.", + ) + parser.add_argument( + "--enable-mla-l1-5-cache", + action="store_true", + help="Enable MLA L1.5 cache in disaggregation paths.", + ) + # Mamba Cache + parser.add_argument( + "--mamba-ssm-dtype", + type=str, + default=ServerArgs.mamba_ssm_dtype, + choices=["float32", "bfloat16"], + help="It is used to tune mamba ssm dtype", + ) + parser.add_argument( + "--mamba-track-interval", + type=int, + default=ServerArgs.mamba_track_interval, + help="The interval to track the mamba state during decode.", + ) + parser.add_argument( + "--max-mamba-cache-size", + type=int, + default=ServerArgs.max_mamba_cache_size, + help="The maximum number of Mamba cache chunks. If unset, the pool size is profiled from available memory.", + ) + parser.add_argument( + "--mamba-full-memory-ratio", + type=float, + default=ServerArgs.mamba_full_memory_ratio, + help="Memory ratio used to split cache budget between Mamba state chunks and full-attention KV cache.", + ) + parser.add_argument( + "--enable-mamba-l2", + action="store_true", + help="Enable host-memory L2 cache for Mamba state slots.", + ) + parser.add_argument( + "--mamba-l2-host-slots", + type=int, + default=ServerArgs.mamba_l2_host_slots, + help="Number of host Mamba L2 slots. If 0, derive from --mamba-l2-host-gb or --mamba-l2-ratio.", + ) + parser.add_argument( + "--mamba-l2-ratio", + type=float, + default=ServerArgs.mamba_l2_ratio, + help="Mamba host L2 slot ratio relative to device Mamba slots when host slots are not explicit.", + ) + parser.add_argument( + "--mamba-l2-layout", + type=str, + choices=["layer_first"], + default=ServerArgs.mamba_l2_layout, + help="Mamba host L2 memory layout.", + ) + parser.add_argument( + "--mamba-l2-io-backend", + type=str, + choices=["direct", "kernel"], + default=ServerArgs.mamba_l2_io_backend, + help="IO backend for Mamba L2 host/device transfers.", + ) + parser.add_argument( + "--mamba-l2-host-gb", + type=int, + default=ServerArgs.mamba_l2_host_gb, + help="Mamba L2 host memory budget in GiB. Overrides --mamba-l2-ratio when host slots are not explicit.", + ) + + parser.add_argument( + "--max-prefill-tokens", + metavar="MAX_PREFILL_TOKENS", + type=int, + default=ServerArgs.max_prefill_tokens, + help=( + "Maximum prefill-token budget used when chunked prefill is " + "disabled. Per-iteration scheduling is controlled by " + "--chunked-prefill-size." + ), + ) + # Other runtime options + parser.add_argument( + "--stream-interval", + type=int, + default=ServerArgs.stream_interval, + help="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", + ) + parser.add_argument( + "--stream-output", + action="store_true", + help="Whether to output as a sequence of disjoint segments.", + ) + parser.add_argument( + "--seed", + metavar="SEED", + type=int, + default=ServerArgs.seed, + help="The random seed.", + ) + parser.add_argument( + "--distributed-timeout-seconds", + metavar="DISTRIBUTED_TIMEOUT_SECONDS", + type=int, + default=ServerArgs.distributed_timeout_seconds, + help="Set timeout for torch.distributed initialization.", + ) + parser.add_argument( + "--download-dir", + type=str, + default=ServerArgs.download_dir, + help="Model download directory for huggingface.", + ) + parser.add_argument( + "--base-gpu-id", + type=int, + default=ServerArgs.base_gpu_id, + help="The base GPU ID to start allocating GPUs from. Useful when running multiple instances on the same machine.", + ) + parser.add_argument( + "--gpu-id-step", + type=int, + default=ServerArgs.gpu_id_step, + help="The delta between consecutive GPU IDs that are used. For example, setting it to 2 will use GPU 0,2,4,...", + ) + + # Logging + parser.add_argument( + "--log-level", + type=str, + default=ServerArgs.log_level, + help="The logging level of all loggers.", + ) + parser.add_argument( + "--log-level-http", + type=str, + default=ServerArgs.log_level_http, + help="The logging level of HTTP server. If not set, reuse --log-level by default.", + ) + parser.add_argument( + "--enable-log-requests", + action=argparse.BooleanOptionalAction, + default=ServerArgs.enable_log_requests, + help="Log metadata, inputs, outputs of all requests. The verbosity is decided by --log-requests-level", + ) + parser.add_argument( + "--log-requests-level", + type=int, + default=0, + help="0: Log metadata. 1. Log metadata and partial input/output. 2. Log every input/output.", + choices=[0, 1, 2], + ) + parser.add_argument( + "--enable-log-request-stats", + action=argparse.BooleanOptionalAction, + default=ServerArgs.enable_log_request_stats, + help=( + "Log a one-line per-request performance summary when each request " + "finishes or aborts: timings (queue/prefill/ttft/total/preemption), " + "token counts (prompt/cache/output), cache-hit rate, decode " + "throughput, and spec-decode acceptance. Measured entirely on the " + "host (no GPU sync), so it adds no engine slowdown." + ), + ) + parser.add_argument( + "--enable-metrics", + action="store_true", + help="Enable log metrics.", + ) + parser.add_argument( + "--metrics-reporters", + action="append", + choices=["prometheus"], + default=["prometheus"], + help="Select metrics reporter(can be specified multiple times)", + ) + + parser.add_argument( + "--app-key", + type=str, + default=ServerArgs.app_key, + help="Set app key of the server", + ) + + parser.add_argument( + "--decode-log-interval", + type=int, + default=ServerArgs.decode_log_interval, + help="The log interval of decode batch.", + ) + # API related + parser.add_argument( + "--api-key", + type=str, + default=ServerArgs.api_key, + help="Set API key of the server. It is also used in the OpenAI API compatible server.", + ) + parser.add_argument( + "--enable-cache-report", + action="store_true", + help="Return number of cached tokens in usage.prompt_tokens_details for each openai request.", + ) + parser.add_argument( + "--kv-events-config", + type=str, + default=ServerArgs.kv_events_config, + help=( + "JSON KV cache event publisher config. Set " + "'enable_kv_cache_events': true and publisher 'zmq' to " + "publish device prefix-cache mutations." + ), + ) + + # Data parallelism + parser.add_argument( + "--data-parallel-size", + metavar="DATA_PARALLEL_SIZE", + type=int, + default=ServerArgs.data_parallel_size, + help="The data parallelism size. If not set, inferred from world_size and attn_tp_size.", + ) + parser.add_argument( + "--load-balance-method", + type=str, + default=ServerArgs.load_balance_method, + help="The load balancing strategy for data parallelism.", + choices=[ + "round_robin", + "shortest_queue", + "minimum_cache_usage", + ], + ) + parser.add_argument( + "--load-watch-interval", + type=float, + default=ServerArgs.load_watch_interval, + help="The interval of load watching in seconds.", + ) + + # Expert parallelism + parser.add_argument( + "--expert-parallel-size", + "--ep-size", + type=int, + default=ServerArgs.ep_size, + help="The expert parallelism size.", + ) + parser.add_argument( + "--init-expert-location", + type=str, + default=ServerArgs.init_expert_location, + help="Initial location of EP experts.", + ) + parser.add_argument( + "--ep-num-redundant-experts", + type=int, + default=ServerArgs.ep_num_redundant_experts, + help="Allocate this number of redundant experts in expert parallel.", + ) + parser.add_argument( + "--ep-dispatch-algorithm", + type=str, + default=ServerArgs.ep_dispatch_algorithm, + help="The algorithm to choose ranks for redundant experts in expert parallel.", + ) + parser.add_argument( + "--eplb-algorithm", + type=str, + default=ServerArgs.eplb_algorithm, + help="Chosen EPLB algorithm", + ) + parser.add_argument( + "--expert-distribution-recorder-mode", + type=str, + default=ServerArgs.expert_distribution_recorder_mode, + help="Mode of expert distribution recorder.", + ) + parser.add_argument( + "--expert-distribution-recorder-buffer-size", + type=int, + default=ServerArgs.expert_distribution_recorder_buffer_size, + help="Circular buffer size of expert distribution recorder. Set to -1 to denote infinite buffer.", + ) + parser.add_argument( + "--enable-expert-distribution-metrics", + action="store_true", + help="Enable logging metrics for expert balancedness", + ) + parser.add_argument( + "--enable-eplb", + action="store_true", + help="Enable EPLB algorithm", + ) + parser.add_argument( + "--moe-backend", + type=str, + default=ServerArgs.moe_backend, + help="MoE runner backend: auto, triton, gluon, flashinfer_trtllm", + ) + parser.add_argument( + "--draft-moe-backend", + type=str, + default=ServerArgs.draft_moe_backend, + help="MoE runner backend for the draft model in speculative decoding. " + "If not set, defaults to --moe-backend.", + ) + parser.add_argument( + "--all2all-backend", + metavar="ALL2ALL_BACKEND", + type=str, + default=ServerArgs.all2all_backend, + help="MoE all-to-all backend: none, deepep, etc.", + ) + parser.add_argument( + "--deepep-mode", + type=str, + choices=["normal", "low_latency", "auto"], + default=ServerArgs.deepep_mode, + help="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.", + ) + parser.add_argument( + "--disable-flashinfer-cutlass-moe-fp4-allgather", + action="store_true", + help="Disable flashinfer cutlass MoE FP4 allgather.", + ) + + # Multi-node distributed serving + parser.add_argument( + "--dist-init-addr", + type=str, + help="The host address for initializing distributed backend (e.g., `192.168.0.2:25000`).", + ) + parser.add_argument( + "--nnodes", type=int, default=ServerArgs.nnodes, help="The number of nodes." + ) + parser.add_argument( + "--node-rank", type=int, default=ServerArgs.node_rank, help="The node rank." + ) + + # Model override args + parser.add_argument( + "--hf-overrides", + metavar="HF_OVERRIDES", + type=str, + help="A dictionary in JSON string format used to override default model configurations.", + default=ServerArgs.hf_overrides, + ) + parser.add_argument( + "--preferred-sampling-params", + type=str, + help="json-formatted sampling settings that will be returned in /get_model_info", + ) + + # Kernel backend + attention_backend_choices = [ + "mha", + "mla", + "fa3", + "fa4", + "triton", + "flashinfer", + "trtllm", + "trtllm_mla", + "flashmla", + "tokenspeed_mla", + "hybrid_linear_attn", + ] + parser.add_argument( + "--attention-backend", + type=str, + choices=attention_backend_choices, + default=ServerArgs.attention_backend, + help="Choose the kernels for attention layers.", + ) + parser.add_argument( + "--drafter-attention-backend", + type=str, + choices=attention_backend_choices, + help="Attention backend for drafter model in speculative decoding. " + "If not specified, uses the same backend as the main model (attention_backend).", + ) + parser.add_argument( + "--sampling-backend", + type=str, + choices=[ + "greedy", + "flashinfer", + "flashinfer_full", + "triton", + "triton_full", + ], + default=ServerArgs.sampling_backend, + help="Sampling backend. " + "'greedy': argmax + verify_chain_greedy, zero sampling-param plumbing. " + "'flashinfer': temperature/top_k/top_p via fused softmax + top_k_top_p_sampling_from_probs; " + "min_p and penalties silently ignored. " + "'triton': temperature/top_k/top_p via MRV2-style logits-to-Gumbel-Max; " + "min_p and penalties silently ignored. " + "'flashinfer_full': adds min_p plus frequency/presence/repetition penalties and logit_bias " + "via the softmax+renorm+min_p kernel sequence. " + "'triton_full': adds min_p plus frequency/presence/repetition penalties and logit_bias " + "with Triton Gumbel-Max for single-step sampling. " + "Allocates a counts[max_req_pool_size, vocab_size] int32 buffer (substantial memory). " + "Finite top_k values must be < 128 or -1.", + ) + parser.add_argument( + "--dp-sampling", + action="store_true", + default=ServerArgs.dp_sampling, + help=( + "Enable Batch-DP spec-verify sampling. Backend selection defaults " + "to auto; override with TOKENSPEED_DP_SAMPLING_BACKEND." + ), + ) + parser.add_argument( + "--dp-sampling-min-bs", + type=int, + default=ServerArgs.dp_sampling_min_bs, + help="Minimum effective decode batch for Batch-DP spec-verify. " + "Defaults to 2 * TP size.", + ) + parser.add_argument( + "--attention-use-fp4-indexer-cache", + "--attention-config.use-fp4-indexer-cache", + "--attention_config.use_fp4_indexer_cache", + type=str_to_bool, + nargs="?", + const=True, + default=ServerArgs.attention_use_fp4_indexer_cache, + help="Use the MXFP4 sparse attention indexer cache layout.", + ) + parser.add_argument( + "--attention-config.use-trtllm-ragged-deepseek-prefill", + "--attention-config.use_trtllm_ragged_deepseek_prefill", + "--attention_config.use_trtllm_ragged_deepseek_prefill", + dest="use_trtllm_ragged_deepseek_prefill", + type=str_to_bool, + nargs="?", + const=True, + default=ServerArgs.use_trtllm_ragged_deepseek_prefill, + help="Use ragged prefill for DeepSeek MLA attention.", + ) + parser.add_argument( + "--deepseek-v4-mega-moe-max-num-tokens", + type=int, + default=ServerArgs.deepseek_v4_mega_moe_max_num_tokens, + help=( + "DeepSeek V4 MegaMoE staging-buffer cap on tokens per forward " + "(0 = derive from chunked-prefill / cuda-graph budgets)." + ), + ) + parser.add_argument( + "--deepseek-v4-indexer-prefill-max-logits-mb", + type=int, + default=ServerArgs.deepseek_v4_indexer_prefill_max_logits_mb, + help=( + "DeepSeek V4 sparse indexer prefill workspace cap (MiB) for the " + "softplus_sqrt logits buffer." + ), + ) + parser.add_argument( + "--deepseek-v4-prefill-chunk-size", + type=int, + default=ServerArgs.deepseek_v4_prefill_chunk_size, + help=( + "Maximum number of requests per DeepSeek V4 FlashMLA prefill " "chunk." + ), + ) + parser.add_argument( + "--grammar-backend", + type=str, + choices=["xgrammar", "none"], + default=ServerArgs.grammar_backend, + help="Grammar backend. 'none' disables grammar-guided decoding entirely ", + ) + parser.add_argument( + "--reasoning-parser", + type=str, + default=ServerArgs.reasoning_parser, + help=( + "Reasoning parser name (e.g. 'minimax', 'kimi_k25'). " + "Used to defer json_schema grammars past the model's " + "reasoning channel." + ), + ) + parser.add_argument( + "--grammar-compile-timeout-secs", + type=float, + default=ServerArgs.grammar_compile_timeout_secs, + help="Per-compile wallclock budget before the request is aborted.", + ) + parser.add_argument( + "--grammar-compile-max-retries", + type=int, + default=ServerArgs.grammar_compile_max_retries, + help="Compile timeouts allowed before a grammar key is permanently rejected.", + ) + parser.add_argument( + "--disable-any-whitespace", + action="store_true", + default=ServerArgs.disable_any_whitespace, + help="Compile xgrammar JSON grammars in tight mode (no arbitrary " + "whitespace between tokens). Mitigates models that wedge into " + "endless whitespace until length cutoff. xgrammar only.", + ) + parser.add_argument( + "--disable-capturable-grammar", + action="store_true", + default=ServerArgs.disable_capturable_grammar, + help="Force the synchronous eager grammar fallback even on CUDA. " + "For parity-testing the captured-grammar path: output should " + "match; throughput will be lower (sync stall every step).", + ) + parser.add_argument( + "--mla-disable-ragged", + action="store_true", + help="Disable the ragged prefill wrapper on MLA kernel backends during EXTEND.", + ) + + # Speculative decoding + parser.add_argument( + "--draft-model-path-use-base", + action="store_true", + help="The path of the draft model weights use the path of the base model", + ) + parser.add_argument( + "--speculative-config", + "--speculative_config", + type=str, + default=ServerArgs.speculative_config, + help="JSON speculative decoding configuration. Supported keys are method, model, and num_speculative_tokens.", + ) + parser.add_argument( + "--speculative-algorithm", + type=str, + choices=["EAGLE3", "MTP", "DFLASH"], + help="Speculative algorithm.", + ) + parser.add_argument( + "--speculative-draft-model-path", + type=str, + help="The path of the draft model weights. This can be a local folder or a Hugging Face repo ID.", + ) + parser.add_argument( + "--speculative-draft-model-quantization", + type=str, + default=ServerArgs.speculative_draft_model_quantization, + help="Quantization method for the draft model. Defaults to 'unquant'.", + ) + parser.add_argument( + "--speculative-num-steps", + type=int, + help="The number of steps sampled from draft model in Speculative Decoding.", + default=ServerArgs.speculative_num_steps, + ) + parser.add_argument( + "--speculative-eagle-topk", + type=int, + help="The number of tokens sampled from the draft model in each speculative step.", + choices=[1], + default=ServerArgs.speculative_eagle_topk, + ) + parser.add_argument( + "--speculative-num-draft-tokens", + type=int, + help="The number of tokens sampled from the draft model in Speculative Decoding.", + default=ServerArgs.speculative_num_draft_tokens, + ) + parser.add_argument( + "--enable-output-logprobs", + action="store_true", + default=ServerArgs.enable_output_logprobs, + help="Enable per-token sampled-token logprobs. OFF by default; enabling extends the captured CUDA-graph footprint. Requests asking for logprobs on a server without this flag receive empty logprobs.", + ) + parser.add_argument( + "--eagle3-layers-to-capture", + type=str, + help="The layers of Eagle3 to capture.", + default=ServerArgs.eagle3_layers_to_capture, + ) + + # Runtime options + parser.add_argument( + "--disable-pdl", + action="store_true", + help="Disable PDL launch.", + ) + prefix_cache_group = parser.add_mutually_exclusive_group() + prefix_cache_group.add_argument( + "--enable-prefix-caching", + action="store_true", + default=ServerArgs.enable_prefix_caching, + help="Enable prefix caching.", + ) + prefix_cache_group.add_argument( + "--no-enable-prefix-caching", + dest="enable_prefix_caching", + action="store_false", + help="Disable prefix caching.", + ) + parser.add_argument( + "--enforce-eager", + action="store_true", + help="Disable CUDA graph.", + ) + parser.add_argument( + "--disable-cuda-graph-padding", + action="store_true", + help="Disable cuda graph when padding is needed. Still uses cuda graph when padding is not needed.", + ) + parser.add_argument( + "--enable-cudagraph-gc", + action="store_true", + help="Enable garbage collection during CUDA graph capture. If disabled (default), GC is frozen during capture to speed up the process.", + ) + parser.add_argument( + "--enable-nccl-nvls", + action="store_true", + help="Enable NCCL NVLS for prefill heavy requests when available.", + ) + parser.add_argument( + "--enable-symm-mem", + action="store_true", + help="Enable NCCL symmetric memory for fast collectives.", + ) + parser.add_argument( + "--disable-custom-all-reduce", + action="store_true", + help="Disable the custom all-reduce kernel and fall back to NCCL.", + ) + parser.add_argument( + "--disable-overlap-schedule", + action="store_true", + help="Disable the overlap scheduler, which overlaps the CPU scheduler with GPU model worker.", + ) + parser.add_argument( + "--disable-tf32", + action="store_true", + help="Disable forcing TF32 on for cuBLAS/cuDNN. By default the server sets " + "NVIDIA_TF32_OVERRIDE=1 and TORCH_ALLOW_TF32_CUBLAS_OVERRIDE=1.", + ) + parser.add_argument( + "--max-cudagraph-capture-size", + metavar="MAX_CUDAGRAPH_CAPTURE_SIZE", + type=int, + default=ServerArgs.max_cudagraph_capture_size, + help="Set the maximum batch size for CUDA graph capture.", + ) + parser.add_argument( + "--cudagraph-capture-sizes", + metavar="CUDAGRAPH_CAPTURE_SIZE", + type=int, + nargs="+", + help="Set the list of batch sizes for CUDA graph capture.", + ) + parser.add_argument( + "--disable-prefill-graph", + action="store_true", + help="Disable cuda graph for prefill.", + ) + parser.add_argument( + "--prefill-graph-max-tokens", + type=int, + default=ServerArgs.prefill_graph_max_tokens, + help="Largest token bucket captured by the breakable prefill CUDA " + "graph. Default (unset) = min(2048, chunked-prefill size); " + "0 disables.", + ) + parser.add_argument( + "--prefill-graph-capture-sizes", + metavar="PREFILL_GRAPH_CAPTURE_SIZE", + type=int, + nargs="+", + help="Explicit list of token-bucket sizes to capture for the " + "breakable prefill graph (like --cudagraph-capture-sizes for " + "decode). Unset: a relative-stride ladder bounding padded compute " + "at ~12.5%% of any size.", + ) + parser.add_argument( + "--enable-nan-detection", + action="store_true", + help="Enable the NaN guard: sanitize non-finite logits before " + "sampling, detect requests whose logits contained NaN (or whose " + "sampled token id escaped the vocab range), and terminate only " + "those requests with a numerical error so corruption cannot " + "spread to the rest of the batch.", + ) + parser.add_argument( + "--enable-nvtx", + action="store_true", + help="Emit NVTX ranges around input_prep / target_forward / " + "sampling / drafter stages for nsys profiling. Off by default " + "(true no-op — no NVTX calls are made). Also enabled by " + "TOKENSPEED_NVTX=1.", + ) + parser.add_argument( + "--enable-p2p-check", + action="store_true", + help="Enable the full GPU P2P access check, otherwise trust the driver's P2P report.", + ) + parser.add_argument( + "--triton-attention-reduce-in-fp32", + action="store_true", + help="Cast the intermediate attention results to fp32 to avoid possible crashes related to fp16." + "This only affects Triton attention kernels.", + ) + parser.add_argument( + "--delete-ckpt-after-loading", + action="store_true", + help="Delete the model checkpoint after loading the model.", + ) + parser.add_argument( + "--weight-loader-prefetch-checkpoints", + action="store_true", + help=( + "Prefetch safetensors checkpoint shards into OS page cache before " + "loading. Local ranks split the shard list to reduce repeated reads " + "from shared filesystems." + ), + ) + parser.add_argument( + "--weight-loader-prefetch-num-threads", + type=int, + default=ServerArgs.weight_loader_prefetch_num_threads, + help="Number of background threads per rank for checkpoint prefetching.", + ) + parser.add_argument( + "--enable-memory-saver", + action="store_true", + help="Allow saving memory using release_memory_occupation and resume_memory_occupation", + ) + parser.add_argument( + "--enable-custom-logit-processor", + action="store_true", + help="Enable users to pass custom logit processors to the server (disabled by default for security)", + ) + # Server warmups + parser.add_argument( + "--skip-server-warmup", + action="store_true", + help="If set, skip warmup.", + ) + parser.add_argument( + "--warmups", + type=str, + required=False, + help="Specify custom warmup functions (csv) to run before server starts eg. --warmups=warmup_name1,warmup_name2 " + "will run the functions `warmup_name1` and `warmup_name2` specified in warmup.py before the server starts listening for requests", + ) + + parser.add_argument( + "--tensor-parallel-size", + "--tp", + type=int, + default=None, + help="Sets tensor parallelism size uniformly (equivalent to --attn-tp-size). " + "Cannot be used together with --attn-tp-size.", + ) + parser.add_argument( + "--enable-expert-parallel", + action="store_true", + help="Enable expert parallelism by automatically setting ep_size to world_size.", + ) + + # Specify different parallel strategies, different combinations correspond to different communication groups and weight partitioning, as well as different communication methods + parser.add_argument( + "--attn-tp-size", + type=int, + default=ServerArgs.attn_tp_size, + help="Specify tp size for attn part", + ) + parser.add_argument( + "--dense-tp-size", + type=int, + default=ServerArgs.dense_tp_size, + help="Specify tp size for dense part, default equals nprocs-per-node, if non dp_attn && combine_dense mode, this parameter will be overridden by attn_tp_size", + ) + parser.add_argument( + "--moe-tp-size", + type=int, + default=ServerArgs.moe_tp_size, + help="Specify tp size for MoE part, default equals nprocs-per-node, if non dp_attn && combine_dense mode, this parameter will be overridden by attn_tp_size", + ) + parser.add_argument( + "--nprocs-per-node", + type=int, + default=ServerArgs.nprocs_per_node, + help="Number of processes to start per node", + ) + parser.add_argument( + "--world-size", + type=int, + default=ServerArgs.world_size, + help="Total number of processes across all nodes.", + ) + parser.add_argument( + "--force-deterministic-rsag", + action="store_true", + help="Enable force deterministic rsag.", + ) + parser.add_argument( + "--disable-sampling-tp-sync", + action="store_true", + help="Skip broadcasting sampler outputs across the attention TP " + "group. Only safe when the sampling kernels are deterministic.", + ) + parser.add_argument( + "--low-latency-max-num-tokens-per-gpu", + type=int, + default=ServerArgs.low_latency_max_num_tokens_per_gpu, + help="Low latency max num tokens per gpu", + ) + + parser.add_argument( + "--mla-chunk-multiplier", + type=int, + default=ServerArgs.mla_chunk_multiplier, + help=( + "Per-iter MLA chunked-prefill chunk capacity multiplier; " + "the actual capacity is chunked_prefill_size * mla_chunk_multiplier." + ), + ) + + # Multimodal + mm_attention_backend_choices = [ + "fa3", + "fa4", + "triton_attn", + "flashinfer_cudnn", + ] + parser.add_argument( + "--mm-attention-backend", + type=str, + choices=mm_attention_backend_choices, + default=ServerArgs.mm_attention_backend, + help="Set multimodal attention backend.", + ) + # Disaggregation + parser.add_argument( + "--disaggregation-mode", + type=str, + default="null", + choices=["null", "prefill", "decode", "encode"], + help='Used for PD/EPD disaggregation. "prefill" for prefill-only server, "decode" for decode-only server, and "encode" for a vision-tower-only server that ships image embeddings to a prefill server. If not specified, it is not disaggregated', + ) + parser.add_argument( + "--comm-fusion-max-num-tokens", + type=int, + default=ServerArgs.comm_fusion_max_num_tokens, + help="Max num tokens for communication fusion workspace", + ) + parser.add_argument( + "--enable-allreduce-fusion", + action="store_true", + help="Enable allreduce fusion for improved decode performance. Auto-enabled on supported single-node TP configurations.", + ) + parser.add_argument( + "--disaggregation-bootstrap-port", + type=int, + default=ServerArgs.disaggregation_bootstrap_port, + help="Bootstrap server port on the prefill server. Default is 8998.", + ) + parser.add_argument( + "--disaggregation-transfer-backend", + type=str, + default=ServerArgs.disaggregation_transfer_backend, + choices=["mooncake", "mooncake_async"], + help="The backend for disaggregation transfer. Default is mooncake.", + ) + parser.add_argument( + "--disaggregation-ib-device", + type=str, + default=ServerArgs.disaggregation_ib_device, + help="The InfiniBand devices for disaggregation transfer, accepts single device (e.g., --disaggregation-ib-device mlx5_0) " + "or multiple comma-separated devices (e.g., --disaggregation-ib-device mlx5_0,mlx5_1). " + "Default is None, which triggers automatic device detection when mooncake backend is enabled.", + ) + parser.add_argument( + "--disaggregation-layerwise-interval", + type=int, + default=ServerArgs.disaggregation_layerwise_interval, + help="The interval of layerwise transfer for disaggregation. Default is 1.", + ) + parser.add_argument( + "--pdlb-url", + type=str, + default=None, + help="The URL of the PD disaggregation load balancer. If set, the prefill/decode server will register with the load balancer.", + ) + + @classmethod + def from_cli_args(cls, args: argparse.Namespace): + args.ep_size = args.expert_parallel_size + + # Resolve model (positional model arg vs --model) + positional_model = getattr(args, "model_path", None) + if positional_model is not None and args.model is not None: + raise ValueError( + "Cannot specify model both as a positional argument and --model. " + "Use one or the other." + ) + if positional_model is not None: + args.model = positional_model + if args.model is None: + raise ValueError( + "Model is required. Provide it as a positional argument " + "(e.g., `tokenspeed serve `) or via --model/--model-path." + ) + + # --tensor-parallel-size → --attn-tp-size + tensor_parallel_size = getattr(args, "tensor_parallel_size", None) + if tensor_parallel_size is not None: + if args.attn_tp_size is not None: + raise ValueError( + "Cannot specify both --tensor-parallel-size and --attn-tp-size. " + "--tensor-parallel-size is an alias for --attn-tp-size." + ) + args.attn_tp_size = tensor_parallel_size + + # Only pass fields that argparse actually produced. Falling back to + # ``None`` for missing attrs would silently clobber dataclass defaults + # for non-CLI-exposed fields (e.g. ``enable_inline_detokenizer``). + attrs = [attr.name for attr in dataclasses.fields(cls)] + return cls( + **{attr: getattr(args, attr) for attr in attrs if hasattr(args, attr)} + ) + + def url(self): + if is_valid_ipv6_address(self.host): + return f"http://[{self.host}]:{self.port}" + return f"http://{self.host}:{self.port}" + + +def prepare_server_args(argv: list[str]) -> ServerArgs: + """ + Prepare the server arguments from the command line arguments. + + Args: + args: The command line arguments. Typically, it should be `sys.argv[1:]`. + + Returns: + The server arguments. + """ + parser = argparse.ArgumentParser(allow_abbrev=False) + ServerArgs.add_cli_args(parser) + raw_args = parser.parse_args(argv) + server_args = ServerArgs.from_cli_args(raw_args) + return server_args + + +ZMQ_TCP_PORT_DELTA = 233 + + +@dataclasses.dataclass +class PortArgs: + # The ipc filename for AsyncLLM to receive BatchTokenIDOut directly + # from the scheduler (zmq). + tokenizer_ipc_name: str + # The ipc filename for scheduler (rank 0) to receive inputs from tokenizer (zmq) + scheduler_input_ipc_name: str + + # The port for nccl initialization (torch.dist) + nccl_port: int + + # The ipc filename for rpc call between Engine and Scheduler + rpc_ipc_name: str + + # The ipc filename for Scheduler to send metrics + metrics_ipc_name: str + + # The ipc filename for Tokenizer and worker tokenizer + tokenizer_worker_ipc_name: str | None + + @staticmethod + def init_new(server_args: ServerArgs, dp_rank: int | None = None) -> "PortArgs": + port = server_args.port + random.randint(100, 1000) + while True: + if is_port_available(port): + break + if port < 60000: + port += 42 + else: + port -= 43 + + # DP attention. Use TCP + port to handle both single-node and multi-node. + if server_args.mapping.nnodes == 1 and server_args.dist_init_addr is None: + # Only use default port fallback when dp_size == 1 + # For dp_size > 1, we need explicit dist_init_addr to avoid port conflicts + if server_args.mapping.has_attn_dp: + raise ValueError( + f"When dp_size > 1 (dp_size={server_args.mapping.attn.dp_size}), you must provide --dist-init-addr. " + f"Example: --dist-init-addr 127.0.0.1:4000" + ) + dist_init_addr = ("127.0.0.1", server_args.port + ZMQ_TCP_PORT_DELTA) + else: + dist_init_addr = server_args.dist_init_addr.split(":") + if len(dist_init_addr) != 2: + raise ValueError( + "please provide --dist-init-addr as host:port of head node" + ) + + dist_init_host, dist_init_port = dist_init_addr + dist_init_port = int(dist_init_port) + + # Scan forward until we find a port cluster where all derived ports are free. + # This handles the case where a previous engine instance left ports in + # TIME_WAIT or its child processes haven't fully terminated yet. + # Note: the port at offset +1 (formerly detokenizer_port) is intentionally + # skipped so the rest of the port layout stays stable for any external + # tooling that indexed off the historical port cluster. + while True: + port_base = dist_init_port + 1 + rpc_port = port_base + 2 + metrics_ipc_port = port_base + 3 + if dp_rank is None: + # TokenizerManager to DataParallelController + scheduler_input_port = port_base + 4 + else: + scheduler_input_port = port_base + 2 + 1 + dp_rank + rpc_ipc_port = scheduler_input_port + 1 + if all( + is_port_available(p) + for p in [ + dist_init_port, + port_base, + rpc_port, + metrics_ipc_port, + scheduler_input_port, + rpc_ipc_port, + ] + ): + break + dist_init_port += 10 + + return PortArgs( + tokenizer_ipc_name=f"tcp://{dist_init_host}:{port_base}", + scheduler_input_ipc_name=f"tcp://{dist_init_host}:{scheduler_input_port}", + nccl_port=port, + rpc_ipc_name=f"tcp://{dist_init_host}:{rpc_port}", + metrics_ipc_name=f"tcp://{dist_init_host}:{metrics_ipc_port}", + tokenizer_worker_ipc_name=None, + ) diff --git a/python/tokenspeed/runtime/utils/text.py b/python/tokenspeed/runtime/utils/text.py new file mode 100644 index 0000000..b1bb5d8 --- /dev/null +++ b/python/tokenspeed/runtime/utils/text.py @@ -0,0 +1,45 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + + +def _is_cjk_char(codepoint: int): + if ( + (0x4E00 <= codepoint <= 0x9FFF) + or (0x3400 <= codepoint <= 0x4DBF) + or (0x20000 <= codepoint <= 0x2A6DF) + or (0x2A700 <= codepoint <= 0x2B73F) + or (0x2B740 <= codepoint <= 0x2B81F) + or (0x2B820 <= codepoint <= 0x2CEAF) + or (0xF900 <= codepoint <= 0xFAFF) + or (0x2F800 <= codepoint <= 0x2FA1F) + ): + return True + + return False + + +def find_printable_text(text: str): + if text.endswith("\n"): + return text + if len(text) > 0 and _is_cjk_char(ord(text[-1])): + return text + if len(text) > 1 and _is_cjk_char(ord(text[-2])): + return text[:-1] + return text[: text.rfind(" ") + 1] diff --git a/python/tokenspeed/runtime/utils/torch_memory_saver_adapter.py b/python/tokenspeed/runtime/utils/torch_memory_saver_adapter.py new file mode 100755 index 0000000..4fa8c25 --- /dev/null +++ b/python/tokenspeed/runtime/utils/torch_memory_saver_adapter.py @@ -0,0 +1,85 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +from abc import ABC +from contextlib import contextmanager + +try: + import torch_memory_saver + + _primary_memory_saver = torch_memory_saver.TorchMemorySaver() +except ImportError: + pass + + +class TorchMemorySaverAdapter(ABC): + @staticmethod + def create(enable: bool): + return ( + _TorchMemorySaverAdapterReal() if enable else _TorchMemorySaverAdapterNoop() + ) + + def configure_subprocess(self): + raise NotImplementedError + + def region(self, tag: str | None = None, enable_cpu_backup: bool = False): + raise NotImplementedError + + def pause(self, tag: str | None = None): + raise NotImplementedError + + def resume(self, tag: str | None = None): + raise NotImplementedError + + +class _TorchMemorySaverAdapterReal(TorchMemorySaverAdapter): + def configure_subprocess(self): + return torch_memory_saver.configure_subprocess() + + def region(self, tag: str | None = None, enable_cpu_backup: bool = False): + # tag defaults to "default" in the library; pass it through explicitly so + # weights vs kv_cache regions can be paused/resumed independently. + # enable_cpu_backup=True offloads (and restores) contents to CPU on + # pause (weights); False discards them (kv_cache). + return _primary_memory_saver.region( + tag=tag or "default", enable_cpu_backup=enable_cpu_backup + ) + + def pause(self, tag: str | None = None): + return _primary_memory_saver.pause(tag=tag) + + def resume(self, tag: str | None = None): + return _primary_memory_saver.resume(tag=tag) + + +class _TorchMemorySaverAdapterNoop(TorchMemorySaverAdapter): + @contextmanager + def configure_subprocess(self): + yield + + @contextmanager + def region(self, tag: str | None = None, enable_cpu_backup: bool = False): + yield + + def pause(self, tag: str | None = None): + pass + + def resume(self, tag: str | None = None): + pass diff --git a/python/tokenspeed/runtime/utils/warmup.py b/python/tokenspeed/runtime/utils/warmup.py new file mode 100755 index 0000000..2b23c89 --- /dev/null +++ b/python/tokenspeed/runtime/utils/warmup.py @@ -0,0 +1,65 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +import numpy as np +import tqdm + +from tokenspeed.runtime.engine.async_llm import AsyncLLM +from tokenspeed.runtime.engine.io_struct import GenerateReqInput +from tokenspeed.runtime.utils import get_colorful_logger + +logger = get_colorful_logger(__name__) + +_warmup_registry = {} + + +def warmup(name: str) -> callable: + def decorator(fn: callable): + _warmup_registry[name] = fn + return fn + + return decorator + + +async def execute_warmups(warmup_names: list[str], tokenizer_manager: AsyncLLM): + for warmup_name in warmup_names: + if warmup_name not in _warmup_registry: + logger.warning("Could not find custom warmup %s", warmup_name) + continue + logger.info("Running warmup %s", warmup_name) + await _warmup_registry[warmup_name](tokenizer_manager) + + +@warmup("voice_chat") +async def voice_chat(tokenizer_manager: AsyncLLM): + # this warms up the fused_moe triton kernels and caches them + # if we don't do this we break real time inference for voice chat + for i in tqdm.trange(1, 512): + size = i * 4 + generate_req_input = GenerateReqInput( + input_ids=(np.random.randint(2**16, size=[size])).tolist(), + sampling_params={ + "max_new_tokens": 30, + "temperature": 0.8, + "stop_token_ids": [1], + "min_p": 0.0, + }, + ) + await tokenizer_manager.generate_request(generate_req_input).__anext__() diff --git a/python/tokenspeed/version.py b/python/tokenspeed/version.py new file mode 100755 index 0000000..1f17bfd --- /dev/null +++ b/python/tokenspeed/version.py @@ -0,0 +1,21 @@ +# Copyright (c) 2026 LightSeek Foundation +# +# 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. + +__version__ = "0.1.0" diff --git a/test/__init__.py b/test/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/test/agentic_benchmark/.gitignore b/test/agentic_benchmark/.gitignore new file mode 100644 index 0000000..ee69b00 --- /dev/null +++ b/test/agentic_benchmark/.gitignore @@ -0,0 +1,3 @@ +build_swe_smith_dataset.py +agentic_dataset.json +outputs/ diff --git a/test/agentic_benchmark/glm5.2/tokenspeed/agentic_bench.sh b/test/agentic_benchmark/glm5.2/tokenspeed/agentic_bench.sh new file mode 100755 index 0000000..2d53328 --- /dev/null +++ b/test/agentic_benchmark/glm5.2/tokenspeed/agentic_bench.sh @@ -0,0 +1,142 @@ +#!/usr/bin/bash + +set -euo pipefail + +# Prepare dataset +EVALSCOPE_COMMIT=acd09b44384d53174768bb1063f675420f76fae9 +pip install "evalscope[perf] @ git+https://github.com/modelscope/evalscope.git@${EVALSCOPE_COMMIT}" + +[ -f build_swe_smith_dataset.py ] || wget https://raw.githubusercontent.com/modelscope/evalscope/${EVALSCOPE_COMMIT}/examples/perf/build_swe_smith_dataset.py \ + -O build_swe_smith_dataset.py + +# Note: Only 94 conversations can be built +[ -f agentic_dataset.json ] || python3 build_swe_smith_dataset.py \ + --model zai-org/GLM-5.2 \ + --first-turn-length 50000 \ + --subsequent-turn-length 800 \ + --min-turns 10 \ + --max-turns 15 \ + --number 128 \ + --output-path agentic_dataset.json \ + --num-workers 32 + +# Sweep configs +CONFIGS=( + attn_tp4_moe_tp4 + attn_tp4_moe_ep4 + attn_tp8_moe_tp8 + attn_tp8_moe_ep8 +) + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SERVER_PID= +SERVER_LOG= + +launch_server() { + local config=$1 + SERVER_LOG=/tmp/tokenspeed_server_${config}.log + setsid ${SCRIPT_DIR}/configs/${config}.sh > "$SERVER_LOG" 2>&1 & + SERVER_PID=$! +} + +wait_for_ready() { + local TIMEOUT=600 + local START=$SECONDS + until curl -sf -o /dev/null http://127.0.0.1:8000/readiness; do + if ! kill -0 "$SERVER_PID" 2>/dev/null; then + echo "Server died early. Last log lines:" >&2 + tail -100 "$SERVER_LOG" >&2 + return 1 + fi + if grep -qE "CUDA out of memory|OutOfMemory|RuntimeError|Killed" "$SERVER_LOG"; then + echo "Server hit a fatal error:" >&2 + tail -100 "$SERVER_LOG" >&2 + return 1 + fi + if (( SECONDS - START > TIMEOUT )); then + echo "Timeout after ${TIMEOUT}s waiting for server" >&2 + return 1 + fi + sleep 5 + done + echo "Server ready after $((SECONDS - START))s" +} + +stop_server() { + if [[ -n "$SERVER_PID" ]] && kill -0 "$SERVER_PID" 2>/dev/null; then + echo "Stopping ts serve (pgid $SERVER_PID)..." + kill -TERM -"$SERVER_PID" 2>/dev/null || true + for _ in {1..20}; do + kill -0 "$SERVER_PID" 2>/dev/null || break + sleep 1 + done + kill -KILL -"$SERVER_PID" 2>/dev/null || true + fi + SERVER_PID= +} + +wait_for_port_free() { + local port=${1:-8000} + local timeout=${2:-90} + local start=$SECONDS + while ! python3 -c "import socket; s=socket.socket(); s.bind(('127.0.0.1', $port)); s.close()" 2>/dev/null; do + if (( SECONDS - start > timeout )); then + echo "Port ${port} still in use after ${timeout}s" >&2 + return 1 + fi + sleep 1 + done +} + +trap stop_server EXIT # safety net for Ctrl-C / errors + +# Preflight: bail out if port 8000 is already in use +wait_for_port_free 8000 + +SWEEP_TS=$(date +%Y%m%d_%H%M%S) +SWEEP_DIR="${SCRIPT_DIR}/outputs/${SWEEP_TS}" +echo "Sweep outputs: ${SWEEP_DIR}" + +for CONFIG in "${CONFIGS[@]}"; do + echo "=== Running $CONFIG ===" + launch_server "$CONFIG" + + if ! wait_for_ready; then + stop_server + exit 1 + fi + + echo "Warmup..." + evalscope perf \ + --model nvidia/GLM-5.2-NVFP4 \ + --url http://127.0.0.1:8000/v1/chat/completions \ + --api openai \ + --dataset swe_smith \ + --dataset-path agentic_dataset.json \ + --max-tokens 500 \ + --multi-turn \ + --number 2 \ + --parallel 2 \ + --extra-args '{"ignore_eos": true}' \ + --dataset-offset 68 \ + --outputs-dir /tmp/outputs + + echo "Benchmark..." + evalscope perf \ + --model nvidia/GLM-5.2-NVFP4 \ + --url http://127.0.0.1:8000/v1/chat/completions \ + --api openai \ + --dataset swe_smith \ + --dataset-path agentic_dataset.json \ + --max-tokens 500 \ + --multi-turn \ + --number 4 8 8 16 32 \ + --parallel 1 2 4 8 16 \ + --extra-args '{"ignore_eos": true}' \ + --name $CONFIG \ + --outputs-dir $SWEEP_DIR \ + --no-timestamp + + stop_server + wait_for_port_free 8000 +done diff --git a/test/agentic_benchmark/glm5.2/tokenspeed/collect_outputs.py b/test/agentic_benchmark/glm5.2/tokenspeed/collect_outputs.py new file mode 100755 index 0000000..dc065fc --- /dev/null +++ b/test/agentic_benchmark/glm5.2/tokenspeed/collect_outputs.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +"""Collect a sweep's per-config benchmark_summary.json files into one CSV.""" + +import argparse +import csv +import json +import re +import sys +from pathlib import Path + +COLUMNS = [ + "config", + "Conc.", + "Latency (tps/user)", + "Throughput (tps/gpu)", + "Approx Cache Hit", + "Decoded Tok/Iter", +] + + +def num_gpus_from_config(config: str) -> int: + m = re.search(r"attn_(?:tp|dp)(\d+)", config) + if not m: + raise ValueError(f"Cannot infer GPU count from config name: {config}") + return int(m.group(1)) + + +def collect(sweep_dir: Path): + rows = [] + for config_dir in sorted(p for p in sweep_dir.iterdir() if p.is_dir()): + config = config_dir.name + n_gpus = num_gpus_from_config(config) + for run_dir in sorted(config_dir.iterdir()): + summary_path = run_dir / "benchmark_summary.json" + if not summary_path.is_file(): + continue + s = json.loads(summary_path.read_text()) + tpot_ms = s["TPOT (ms)"] + decode_tps_user = 1000.0 / tpot_ms if tpot_ms else 0.0 + tps_gpu = s["Total Throughput (tok/s)"] / n_gpus + rows.append( + { + "config": config, + "Conc.": s["Concurrency"], + "Latency (tps/user)": round(decode_tps_user, 2), + "Throughput (tps/gpu)": round(tps_gpu, 2), + "Approx Cache Hit": round(s["KV Cache Hit Rate (%)"], 2), + "Decoded Tok/Iter": round(s["Decoded Tok/Iter"], 4), + } + ) + rows.sort(key=lambda r: (r["config"], r["Conc."])) + return rows + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("sweep_dir", type=Path, help="e.g. outputs/20260505_152734") + ap.add_argument( + "-o", + "--output", + type=Path, + default=None, + help="CSV output path (default: stdout)", + ) + args = ap.parse_args() + + if not args.sweep_dir.is_dir(): + sys.exit(f"Not a directory: {args.sweep_dir}") + + rows = collect(args.sweep_dir) + out = args.output.open("w", newline="") if args.output else sys.stdout + try: + w = csv.DictWriter(out, fieldnames=COLUMNS) + w.writeheader() + w.writerows(rows) + finally: + if args.output: + out.close() + + +if __name__ == "__main__": + main() diff --git a/test/agentic_benchmark/glm5.2/tokenspeed/configs/attn_tp4_moe_ep4.sh b/test/agentic_benchmark/glm5.2/tokenspeed/configs/attn_tp4_moe_ep4.sh new file mode 100755 index 0000000..69d766d --- /dev/null +++ b/test/agentic_benchmark/glm5.2/tokenspeed/configs/attn_tp4_moe_ep4.sh @@ -0,0 +1,26 @@ +#!/usr/bin/bash + +set -euo pipefail + +exec ts serve \ + --model nvidia/GLM-5.2-NVFP4 \ + --attn-tp-size 4 \ + --ep-size 4 \ + --max-model-len 80000 \ + --max-num-seqs 16 \ + --max-prefill-tokens 8192 \ + --chunked-prefill-size 8192 \ + --gpu-memory-utilization 0.95 \ + --disable-cuda-graph-padding \ + --trust-remote-code \ + --moe-backend flashinfer_trtllm \ + --quantization nvfp4 \ + --kv-cache-dtype fp8 \ + --speculative-algorithm MTP \ + --speculative-num-steps 3 \ + --speculative-eagle-topk 1 \ + --speculative-num-draft-tokens 4 \ + --speculative-draft-model-quantization unquant \ + --enable-cache-report \ + --host 127.0.0.1 \ + --port 8000 diff --git a/test/agentic_benchmark/glm5.2/tokenspeed/configs/attn_tp4_moe_tp4.sh b/test/agentic_benchmark/glm5.2/tokenspeed/configs/attn_tp4_moe_tp4.sh new file mode 100755 index 0000000..c388183 --- /dev/null +++ b/test/agentic_benchmark/glm5.2/tokenspeed/configs/attn_tp4_moe_tp4.sh @@ -0,0 +1,26 @@ +#!/usr/bin/bash + +set -euo pipefail + +exec ts serve \ + --model nvidia/GLM-5.2-NVFP4 \ + --attn-tp-size 4 \ + --moe-tp-size 4 \ + --max-model-len 80000 \ + --max-num-seqs 16 \ + --max-prefill-tokens 8192 \ + --chunked-prefill-size 8192 \ + --gpu-memory-utilization 0.95 \ + --disable-cuda-graph-padding \ + --trust-remote-code \ + --moe-backend flashinfer_trtllm \ + --quantization nvfp4 \ + --kv-cache-dtype fp8 \ + --speculative-algorithm MTP \ + --speculative-num-steps 3 \ + --speculative-eagle-topk 1 \ + --speculative-num-draft-tokens 4 \ + --speculative-draft-model-quantization unquant \ + --enable-cache-report \ + --host 127.0.0.1 \ + --port 8000 diff --git a/test/agentic_benchmark/glm5.2/tokenspeed/configs/attn_tp8_moe_ep8.sh b/test/agentic_benchmark/glm5.2/tokenspeed/configs/attn_tp8_moe_ep8.sh new file mode 100755 index 0000000..56460fd --- /dev/null +++ b/test/agentic_benchmark/glm5.2/tokenspeed/configs/attn_tp8_moe_ep8.sh @@ -0,0 +1,26 @@ +#!/usr/bin/bash + +set -euo pipefail + +exec ts serve \ + --model nvidia/GLM-5.2-NVFP4 \ + --attn-tp-size 8 \ + --ep-size 8 \ + --max-model-len 80000 \ + --max-num-seqs 16 \ + --max-prefill-tokens 8192 \ + --chunked-prefill-size 8192 \ + --gpu-memory-utilization 0.95 \ + --disable-cuda-graph-padding \ + --trust-remote-code \ + --moe-backend flashinfer_trtllm \ + --quantization nvfp4 \ + --kv-cache-dtype fp8 \ + --speculative-algorithm MTP \ + --speculative-num-steps 3 \ + --speculative-eagle-topk 1 \ + --speculative-num-draft-tokens 4 \ + --speculative-draft-model-quantization unquant \ + --enable-cache-report \ + --host 127.0.0.1 \ + --port 8000 diff --git a/test/agentic_benchmark/glm5.2/tokenspeed/configs/attn_tp8_moe_tp8.sh b/test/agentic_benchmark/glm5.2/tokenspeed/configs/attn_tp8_moe_tp8.sh new file mode 100755 index 0000000..8c0de3c --- /dev/null +++ b/test/agentic_benchmark/glm5.2/tokenspeed/configs/attn_tp8_moe_tp8.sh @@ -0,0 +1,26 @@ +#!/usr/bin/bash + +set -euo pipefail + +exec ts serve \ + --model nvidia/GLM-5.2-NVFP4 \ + --attn-tp-size 8 \ + --moe-tp-size 8 \ + --max-model-len 80000 \ + --max-num-seqs 16 \ + --max-prefill-tokens 8192 \ + --chunked-prefill-size 8192 \ + --gpu-memory-utilization 0.95 \ + --disable-cuda-graph-padding \ + --trust-remote-code \ + --moe-backend flashinfer_trtllm \ + --quantization nvfp4 \ + --kv-cache-dtype fp8 \ + --speculative-algorithm MTP \ + --speculative-num-steps 3 \ + --speculative-eagle-topk 1 \ + --speculative-num-draft-tokens 4 \ + --speculative-draft-model-quantization unquant \ + --enable-cache-report \ + --host 127.0.0.1 \ + --port 8000 diff --git a/test/agentic_benchmark/glm5.2/trtllm/agentic_bench.sh b/test/agentic_benchmark/glm5.2/trtllm/agentic_bench.sh new file mode 100755 index 0000000..fd0a7b8 --- /dev/null +++ b/test/agentic_benchmark/glm5.2/trtllm/agentic_bench.sh @@ -0,0 +1,152 @@ +#!/usr/bin/bash + +# Tested on TensorRT-LLM commit: 281acfd + +set -euo pipefail + +# Prepare dataset +EVALSCOPE_COMMIT=acd09b44384d53174768bb1063f675420f76fae9 +pip install "evalscope[perf] @ git+https://github.com/modelscope/evalscope.git@${EVALSCOPE_COMMIT}" + +[ -f build_swe_smith_dataset.py ] || wget https://raw.githubusercontent.com/modelscope/evalscope/${EVALSCOPE_COMMIT}/examples/perf/build_swe_smith_dataset.py \ + -O build_swe_smith_dataset.py + +# Note: Only 94 conversations can be built +[ -f agentic_dataset.json ] || python3 build_swe_smith_dataset.py \ + --model zai-org/GLM-5.2 \ + --first-turn-length 50000 \ + --subsequent-turn-length 800 \ + --min-turns 10 \ + --max-turns 15 \ + --number 128 \ + --output-path agentic_dataset.json \ + --num-workers 32 + +# Sweep configs +CONFIGS=( + attn_tp4_moe_tp4 + attn_tp4_moe_ep4 + attn_tp8_moe_tp8 + attn_tp8_moe_ep8 +) + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SERVER_PID= +SERVER_LOG= + +launch_server() { + local config=$1 + SERVER_LOG=/tmp/trtllm_server_${config}.log + setsid trtllm-serve nvidia/GLM-5.2-NVFP4 \ + --max_num_tokens 8192 \ + --max_seq_len 80000 \ + --enable_chunked_prefill \ + --trust_remote_code \ + --config "${SCRIPT_DIR}/configs/${config}.yaml" \ + --host 127.0.0.1 \ + --port 8001 \ + > "$SERVER_LOG" 2>&1 & + SERVER_PID=$! +} + +wait_for_ready() { + local TIMEOUT=600 + local START=$SECONDS + until curl -sf -o /dev/null http://127.0.0.1:8001/health; do + if ! kill -0 "$SERVER_PID" 2>/dev/null; then + echo "Server died early. Last log lines:" >&2 + tail -100 "$SERVER_LOG" >&2 + return 1 + fi + if grep -qE "CUDA out of memory|OutOfMemory|RuntimeError|Killed" "$SERVER_LOG"; then + echo "Server hit a fatal error:" >&2 + tail -100 "$SERVER_LOG" >&2 + return 1 + fi + if (( SECONDS - START > TIMEOUT )); then + echo "Timeout after ${TIMEOUT}s waiting for server" >&2 + return 1 + fi + sleep 5 + done + echo "Server ready after $((SECONDS - START))s" +} + +stop_server() { + if [[ -n "$SERVER_PID" ]] && kill -0 "$SERVER_PID" 2>/dev/null; then + echo "Stopping trtllm-serve (pgid $SERVER_PID)..." + kill -TERM -"$SERVER_PID" 2>/dev/null || true + for _ in {1..20}; do + kill -0 "$SERVER_PID" 2>/dev/null || break + sleep 1 + done + kill -KILL -"$SERVER_PID" 2>/dev/null || true + fi + SERVER_PID= +} + +wait_for_port_free() { + local port=${1:-8001} + local timeout=${2:-90} + local start=$SECONDS + while ! python3 -c "import socket; s=socket.socket(); s.bind(('127.0.0.1', $port)); s.close()" 2>/dev/null; do + if (( SECONDS - start > timeout )); then + echo "Port ${port} still in use after ${timeout}s" >&2 + return 1 + fi + sleep 1 + done +} + +trap stop_server EXIT # safety net for Ctrl-C / errors + +# Preflight: bail out if port 8001 is already in use +wait_for_port_free 8001 + +SWEEP_TS=$(date +%Y%m%d_%H%M%S) +SWEEP_DIR="${SCRIPT_DIR}/outputs/${SWEEP_TS}" +echo "Sweep outputs: ${SWEEP_DIR}" + +for CONFIG in "${CONFIGS[@]}"; do + echo "=== Running $CONFIG ===" + launch_server "$CONFIG" + + if ! wait_for_ready; then + stop_server + exit 1 + fi + + echo "Warmup..." + evalscope perf \ + --model nvidia/GLM-5.2-NVFP4 \ + --url http://127.0.0.1:8001/v1/chat/completions \ + --api openai \ + --dataset swe_smith \ + --dataset-path agentic_dataset.json \ + --max-tokens 500 \ + --multi-turn \ + --number 2 \ + --parallel 2 \ + --extra-args '{"ignore_eos": true}' \ + --dataset-offset 68 \ + --outputs-dir /tmp/outputs + + echo "Benchmark..." + evalscope perf \ + --model nvidia/GLM-5.2-NVFP4 \ + --url http://127.0.0.1:8001/v1/chat/completions \ + --api openai \ + --dataset swe_smith \ + --dataset-path agentic_dataset.json \ + --max-tokens 500 \ + --multi-turn \ + --number 4 8 8 16 32 \ + --parallel 1 2 4 8 16 \ + --extra-args '{"ignore_eos": true}' \ + --name $CONFIG \ + --outputs-dir $SWEEP_DIR \ + --no-timestamp + + stop_server + wait_for_port_free 8001 +done diff --git a/test/agentic_benchmark/glm5.2/trtllm/collect_outputs.py b/test/agentic_benchmark/glm5.2/trtllm/collect_outputs.py new file mode 100755 index 0000000..dc065fc --- /dev/null +++ b/test/agentic_benchmark/glm5.2/trtllm/collect_outputs.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +"""Collect a sweep's per-config benchmark_summary.json files into one CSV.""" + +import argparse +import csv +import json +import re +import sys +from pathlib import Path + +COLUMNS = [ + "config", + "Conc.", + "Latency (tps/user)", + "Throughput (tps/gpu)", + "Approx Cache Hit", + "Decoded Tok/Iter", +] + + +def num_gpus_from_config(config: str) -> int: + m = re.search(r"attn_(?:tp|dp)(\d+)", config) + if not m: + raise ValueError(f"Cannot infer GPU count from config name: {config}") + return int(m.group(1)) + + +def collect(sweep_dir: Path): + rows = [] + for config_dir in sorted(p for p in sweep_dir.iterdir() if p.is_dir()): + config = config_dir.name + n_gpus = num_gpus_from_config(config) + for run_dir in sorted(config_dir.iterdir()): + summary_path = run_dir / "benchmark_summary.json" + if not summary_path.is_file(): + continue + s = json.loads(summary_path.read_text()) + tpot_ms = s["TPOT (ms)"] + decode_tps_user = 1000.0 / tpot_ms if tpot_ms else 0.0 + tps_gpu = s["Total Throughput (tok/s)"] / n_gpus + rows.append( + { + "config": config, + "Conc.": s["Concurrency"], + "Latency (tps/user)": round(decode_tps_user, 2), + "Throughput (tps/gpu)": round(tps_gpu, 2), + "Approx Cache Hit": round(s["KV Cache Hit Rate (%)"], 2), + "Decoded Tok/Iter": round(s["Decoded Tok/Iter"], 4), + } + ) + rows.sort(key=lambda r: (r["config"], r["Conc."])) + return rows + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("sweep_dir", type=Path, help="e.g. outputs/20260505_152734") + ap.add_argument( + "-o", + "--output", + type=Path, + default=None, + help="CSV output path (default: stdout)", + ) + args = ap.parse_args() + + if not args.sweep_dir.is_dir(): + sys.exit(f"Not a directory: {args.sweep_dir}") + + rows = collect(args.sweep_dir) + out = args.output.open("w", newline="") if args.output else sys.stdout + try: + w = csv.DictWriter(out, fieldnames=COLUMNS) + w.writeheader() + w.writerows(rows) + finally: + if args.output: + out.close() + + +if __name__ == "__main__": + main() diff --git a/test/agentic_benchmark/glm5.2/trtllm/configs/attn_tp4_moe_ep4.yaml b/test/agentic_benchmark/glm5.2/trtllm/configs/attn_tp4_moe_ep4.yaml new file mode 100644 index 0000000..ef882e1 --- /dev/null +++ b/test/agentic_benchmark/glm5.2/trtllm/configs/attn_tp4_moe_ep4.yaml @@ -0,0 +1,13 @@ +tensor_parallel_size: 4 +enable_attention_dp: false +moe_expert_parallel_size: 4 +max_batch_size: 16 +cuda_graph_config: + batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16] +kv_cache_config: + dtype: fp8 +moe_config: + backend: TRTLLM +speculative_config: + decoding_type: MTP + max_draft_len: 3 diff --git a/test/agentic_benchmark/glm5.2/trtllm/configs/attn_tp4_moe_tp4.yaml b/test/agentic_benchmark/glm5.2/trtllm/configs/attn_tp4_moe_tp4.yaml new file mode 100644 index 0000000..bfec2bb --- /dev/null +++ b/test/agentic_benchmark/glm5.2/trtllm/configs/attn_tp4_moe_tp4.yaml @@ -0,0 +1,13 @@ +tensor_parallel_size: 4 +enable_attention_dp: false +moe_expert_parallel_size: 1 +max_batch_size: 16 +cuda_graph_config: + batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16] +kv_cache_config: + dtype: fp8 +moe_config: + backend: TRTLLM +speculative_config: + decoding_type: MTP + max_draft_len: 3 diff --git a/test/agentic_benchmark/glm5.2/trtllm/configs/attn_tp8_moe_ep8.yaml b/test/agentic_benchmark/glm5.2/trtllm/configs/attn_tp8_moe_ep8.yaml new file mode 100644 index 0000000..5213400 --- /dev/null +++ b/test/agentic_benchmark/glm5.2/trtllm/configs/attn_tp8_moe_ep8.yaml @@ -0,0 +1,13 @@ +tensor_parallel_size: 8 +enable_attention_dp: false +moe_expert_parallel_size: 8 +max_batch_size: 16 +cuda_graph_config: + batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16] +kv_cache_config: + dtype: fp8 +moe_config: + backend: TRTLLM +speculative_config: + decoding_type: MTP + max_draft_len: 3 diff --git a/test/agentic_benchmark/glm5.2/trtllm/configs/attn_tp8_moe_tp8.yaml b/test/agentic_benchmark/glm5.2/trtllm/configs/attn_tp8_moe_tp8.yaml new file mode 100644 index 0000000..f81669d --- /dev/null +++ b/test/agentic_benchmark/glm5.2/trtllm/configs/attn_tp8_moe_tp8.yaml @@ -0,0 +1,13 @@ +tensor_parallel_size: 8 +enable_attention_dp: false +moe_expert_parallel_size: 1 +max_batch_size: 16 +cuda_graph_config: + batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16] +kv_cache_config: + dtype: fp8 +moe_config: + backend: TRTLLM +speculative_config: + decoding_type: MTP + max_draft_len: 3 diff --git a/test/agentic_benchmark/kimi_k2.5/sglang/agentic_bench.sh b/test/agentic_benchmark/kimi_k2.5/sglang/agentic_bench.sh new file mode 100755 index 0000000..7430820 --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/sglang/agentic_bench.sh @@ -0,0 +1,142 @@ +#!/usr/bin/bash + +# Tested on SGLang tag: 0.5.12.post1 + +set -euo pipefail + +# Prepare dataset +EVALSCOPE_COMMIT=acd09b44384d53174768bb1063f675420f76fae9 +pip install "evalscope[perf] @ git+https://github.com/modelscope/evalscope.git@${EVALSCOPE_COMMIT}" + +[ -f build_swe_smith_dataset.py ] || wget https://raw.githubusercontent.com/modelscope/evalscope/${EVALSCOPE_COMMIT}/examples/perf/build_swe_smith_dataset.py \ + -O build_swe_smith_dataset.py + +# Note: Only 71 conversations can be built +[ -f agentic_dataset.json ] || python3 build_swe_smith_dataset.py \ + --model moonshotai/Kimi-K2.5 \ + --first-turn-length 50000 \ + --subsequent-turn-length 800 \ + --min-turns 10 \ + --max-turns 15 \ + --number 128 \ + --output-path agentic_dataset.json \ + --num-workers 32 + +# Sweep configs +CONFIGS=( + attn_tp4_moe_tp4 + attn_tp4_moe_ep4 +) + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SERVER_PID= +SERVER_LOG= + +launch_server() { + local config=$1 + SERVER_LOG=/tmp/sglang_server_${config}.log + setsid ${SCRIPT_DIR}/configs/${config}.sh > "$SERVER_LOG" 2>&1 & + SERVER_PID=$! +} + +wait_for_ready() { + local TIMEOUT=1200 + local START=$SECONDS + until curl -sf -o /dev/null http://127.0.0.1:8003/health; do + if ! kill -0 "$SERVER_PID" 2>/dev/null; then + echo "Server died early. Last log lines:" >&2 + tail -100 "$SERVER_LOG" >&2 + return 1 + fi + if grep -qE "CUDA out of memory|OutOfMemory|RuntimeError|Killed" "$SERVER_LOG"; then + echo "Server hit a fatal error:" >&2 + tail -100 "$SERVER_LOG" >&2 + return 1 + fi + if (( SECONDS - START > TIMEOUT )); then + echo "Timeout after ${TIMEOUT}s waiting for server" >&2 + return 1 + fi + sleep 5 + done + echo "Server ready after $((SECONDS - START))s" +} + +stop_server() { + if [[ -n "$SERVER_PID" ]] && kill -0 "$SERVER_PID" 2>/dev/null; then + echo "Stopping sglang serve (pgid $SERVER_PID)..." + kill -TERM -"$SERVER_PID" 2>/dev/null || true + for _ in {1..20}; do + kill -0 "$SERVER_PID" 2>/dev/null || break + sleep 1 + done + kill -KILL -"$SERVER_PID" 2>/dev/null || true + fi + SERVER_PID= +} + +wait_for_port_free() { + local port=${1:-8003} + local timeout=${2:-90} + local start=$SECONDS + while ! python3 -c "import socket; s=socket.socket(); s.bind(('127.0.0.1', $port)); s.close()" 2>/dev/null; do + if (( SECONDS - start > timeout )); then + echo "Port ${port} still in use after ${timeout}s" >&2 + return 1 + fi + sleep 1 + done +} + +trap stop_server EXIT # safety net for Ctrl-C / errors + +# Preflight: bail out if port 8003 is already in use +wait_for_port_free 8003 + +SWEEP_TS=$(date +%Y%m%d_%H%M%S) +SWEEP_DIR="${SCRIPT_DIR}/outputs/${SWEEP_TS}" +echo "Sweep outputs: ${SWEEP_DIR}" + +for CONFIG in "${CONFIGS[@]}"; do + echo "=== Running $CONFIG ===" + launch_server "$CONFIG" + + if ! wait_for_ready; then + stop_server + exit 1 + fi + + echo "Warmup..." + evalscope perf \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --url http://127.0.0.1:8003/v1/chat/completions \ + --api openai \ + --dataset swe_smith \ + --dataset-path agentic_dataset.json \ + --max-tokens 500 \ + --multi-turn \ + --number 2 \ + --parallel 2 \ + --extra-args '{"ignore_eos": true}' \ + --dataset-offset 68 \ + --outputs-dir /tmp/outputs + + echo "Benchmark..." + evalscope perf \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --url http://127.0.0.1:8003/v1/chat/completions \ + --api openai \ + --dataset swe_smith \ + --dataset-path agentic_dataset.json \ + --max-tokens 500 \ + --multi-turn \ + --number 4 8 8 16 32 \ + --parallel 1 2 4 8 16 \ + --extra-args '{"ignore_eos": true}' \ + --name $CONFIG \ + --outputs-dir $SWEEP_DIR \ + --no-timestamp + + stop_server + wait_for_port_free 8003 +done diff --git a/test/agentic_benchmark/kimi_k2.5/sglang/collect_outputs.py b/test/agentic_benchmark/kimi_k2.5/sglang/collect_outputs.py new file mode 100755 index 0000000..dc065fc --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/sglang/collect_outputs.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +"""Collect a sweep's per-config benchmark_summary.json files into one CSV.""" + +import argparse +import csv +import json +import re +import sys +from pathlib import Path + +COLUMNS = [ + "config", + "Conc.", + "Latency (tps/user)", + "Throughput (tps/gpu)", + "Approx Cache Hit", + "Decoded Tok/Iter", +] + + +def num_gpus_from_config(config: str) -> int: + m = re.search(r"attn_(?:tp|dp)(\d+)", config) + if not m: + raise ValueError(f"Cannot infer GPU count from config name: {config}") + return int(m.group(1)) + + +def collect(sweep_dir: Path): + rows = [] + for config_dir in sorted(p for p in sweep_dir.iterdir() if p.is_dir()): + config = config_dir.name + n_gpus = num_gpus_from_config(config) + for run_dir in sorted(config_dir.iterdir()): + summary_path = run_dir / "benchmark_summary.json" + if not summary_path.is_file(): + continue + s = json.loads(summary_path.read_text()) + tpot_ms = s["TPOT (ms)"] + decode_tps_user = 1000.0 / tpot_ms if tpot_ms else 0.0 + tps_gpu = s["Total Throughput (tok/s)"] / n_gpus + rows.append( + { + "config": config, + "Conc.": s["Concurrency"], + "Latency (tps/user)": round(decode_tps_user, 2), + "Throughput (tps/gpu)": round(tps_gpu, 2), + "Approx Cache Hit": round(s["KV Cache Hit Rate (%)"], 2), + "Decoded Tok/Iter": round(s["Decoded Tok/Iter"], 4), + } + ) + rows.sort(key=lambda r: (r["config"], r["Conc."])) + return rows + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("sweep_dir", type=Path, help="e.g. outputs/20260505_152734") + ap.add_argument( + "-o", + "--output", + type=Path, + default=None, + help="CSV output path (default: stdout)", + ) + args = ap.parse_args() + + if not args.sweep_dir.is_dir(): + sys.exit(f"Not a directory: {args.sweep_dir}") + + rows = collect(args.sweep_dir) + out = args.output.open("w", newline="") if args.output else sys.stdout + try: + w = csv.DictWriter(out, fieldnames=COLUMNS) + w.writeheader() + w.writerows(rows) + finally: + if args.output: + out.close() + + +if __name__ == "__main__": + main() diff --git a/test/agentic_benchmark/kimi_k2.5/sglang/configs/attn_tp4_moe_ep4.sh b/test/agentic_benchmark/kimi_k2.5/sglang/configs/attn_tp4_moe_ep4.sh new file mode 100755 index 0000000..93d6622 --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/sglang/configs/attn_tp4_moe_ep4.sh @@ -0,0 +1,25 @@ +#!/usr/bin/bash + +set -euo pipefail + +exec sglang serve \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --tp-size 4 \ + --ep-size 4 \ + --context-length 80000 \ + --max-running-requests 16 \ + --max-prefill-tokens 8192 \ + --chunked-prefill-size 8192 \ + --mem-fraction-static 0.9 \ + --trust-remote-code \ + --attention-backend trtllm_mla \ + --moe-runner-backend flashinfer_trtllm \ + --kv-cache-dtype fp8_e4m3 \ + --cuda-graph-max-bs 16 \ + --speculative-algorithm EAGLE3 \ + --speculative-draft-model-path lightseekorg/kimi-k2.5-eagle3-mla \ + --speculative-num-steps 3 \ + --speculative-eagle-topk 1 \ + --speculative-num-draft-tokens 4 \ + --host 127.0.0.1 \ + --port 8003 diff --git a/test/agentic_benchmark/kimi_k2.5/sglang/configs/attn_tp4_moe_tp4.sh b/test/agentic_benchmark/kimi_k2.5/sglang/configs/attn_tp4_moe_tp4.sh new file mode 100755 index 0000000..cc07db8 --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/sglang/configs/attn_tp4_moe_tp4.sh @@ -0,0 +1,24 @@ +#!/usr/bin/bash + +set -euo pipefail + +exec sglang serve \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --tp-size 4 \ + --context-length 80000 \ + --max-running-requests 16 \ + --max-prefill-tokens 8192 \ + --chunked-prefill-size 8192 \ + --mem-fraction-static 0.9 \ + --trust-remote-code \ + --attention-backend trtllm_mla \ + --moe-runner-backend flashinfer_trtllm \ + --kv-cache-dtype fp8_e4m3 \ + --cuda-graph-max-bs 16 \ + --speculative-algorithm EAGLE3 \ + --speculative-draft-model-path lightseekorg/kimi-k2.5-eagle3-mla \ + --speculative-num-steps 3 \ + --speculative-eagle-topk 1 \ + --speculative-num-draft-tokens 4 \ + --host 127.0.0.1 \ + --port 8003 diff --git a/test/agentic_benchmark/kimi_k2.5/tokenspeed/README.md b/test/agentic_benchmark/kimi_k2.5/tokenspeed/README.md new file mode 100644 index 0000000..6886fb4 --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/tokenspeed/README.md @@ -0,0 +1,62 @@ +# Agentic Benchmark — TokenSpeed + +Sweep `ts serve` against an agentic, multi-turn workload (SWE-Smith) at a +fixed set of attention/MoE parallelism layouts and report per-config throughput, +latency, and KV-cache hit rate. + +Server listens on port **8000**. + +## Layout + +``` +agentic_bench.sh # main sweep: dataset prep -> for config in CONFIGS: launch, wait, bench, kill +configs/ # one shell script per parallelism layout (each `exec`s ts serve) +collect_outputs.py # parse a sweep into a flat CSV +outputs///parallel_

_number_/ # per-run evalscope artifacts +``` + +## Run a sweep + +```bash +cd test/agentic_benchmark/tokenspeed +./agentic_bench.sh +``` + +The script (1) installs evalscope at the pinned commit, (2) builds the SWE-Smith +multi-turn dataset, (3) iterates each config in `CONFIGS=()`: launch server, poll +`/readiness` until ready, run `evalscope perf`, kill server, wait for the port to be +free, repeat. Aborts the whole sweep on the first failure (`set -e`). + +To narrow the matrix, comment out entries in the `CONFIGS=()` array. + +## Configs + +Each `configs/*.sh` `exec`s `ts serve` with the full flag set for one +layout. Key flags: + +- `--attn-tp-size`, `--moe-tp-size` *or* `--ep-size` +- `--max-num-seqs`, `--max-prefill-tokens`, `--chunked-prefill-size` +- `--quantization nvfp4`, `--kv-cache-dtype fp8` +- `--moe-backend flashinfer_trtllm`, `--attention-backend tokenspeed_mla` +- Eagle3 spec-dec: `--speculative-algorithm EAGLE3 --speculative-num-steps 3 + --speculative-eagle-topk 1 --speculative-num-draft-tokens 4` + +Naming: `attn__moe_` where `X ∈ {tp4,tp8,dp8}` and `Y ∈ {tp4,tp8,ep4,ep8}`. +World size = the number after `attn_(tp|dp)`. + +To verify the parallelism actually applied, grep the server log: +```bash +grep -A6 "Parallelism configuration" /tmp/tokenspeed_server_.log +``` + +## Collect results + +```bash +python3 collect_outputs.py outputs/ -o sweep.csv +``` + +Emits one row per (config, concurrency) with `Conc.`, `Latency (tps/user)`, +`Throughput (tps/gpu)`, `Approx Cache Hit`, `Decoded Tok/Iter`. `tps/gpu` divides +the system-wide `Total Throughput (tok/s)` by the GPU count inferred from the +config name; the other metrics come straight from `benchmark_summary.json` +(same numbers as evalscope's `performance_summary.txt` Request Metrics table). diff --git a/test/agentic_benchmark/kimi_k2.5/tokenspeed/agentic_bench.sh b/test/agentic_benchmark/kimi_k2.5/tokenspeed/agentic_bench.sh new file mode 100755 index 0000000..3e573a4 --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/tokenspeed/agentic_bench.sh @@ -0,0 +1,142 @@ +#!/usr/bin/bash + +set -euo pipefail + +# Prepare dataset +EVALSCOPE_COMMIT=acd09b44384d53174768bb1063f675420f76fae9 +pip install "evalscope[perf] @ git+https://github.com/modelscope/evalscope.git@${EVALSCOPE_COMMIT}" + +[ -f build_swe_smith_dataset.py ] || wget https://raw.githubusercontent.com/modelscope/evalscope/${EVALSCOPE_COMMIT}/examples/perf/build_swe_smith_dataset.py \ + -O build_swe_smith_dataset.py + +# Note: Only 71 conversations can be built +[ -f agentic_dataset.json ] || python3 build_swe_smith_dataset.py \ + --model moonshotai/Kimi-K2.5 \ + --first-turn-length 50000 \ + --subsequent-turn-length 800 \ + --min-turns 10 \ + --max-turns 15 \ + --number 128 \ + --output-path agentic_dataset.json \ + --num-workers 32 + +# Sweep configs +CONFIGS=( + attn_tp4_moe_tp4 + attn_tp4_moe_ep4 + attn_tp8_moe_tp8 + attn_tp8_moe_ep8 +) + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SERVER_PID= +SERVER_LOG= + +launch_server() { + local config=$1 + SERVER_LOG=/tmp/tokenspeed_server_${config}.log + setsid ${SCRIPT_DIR}/configs/${config}.sh > "$SERVER_LOG" 2>&1 & + SERVER_PID=$! +} + +wait_for_ready() { + local TIMEOUT=600 + local START=$SECONDS + until curl -sf -o /dev/null http://127.0.0.1:8000/readiness; do + if ! kill -0 "$SERVER_PID" 2>/dev/null; then + echo "Server died early. Last log lines:" >&2 + tail -100 "$SERVER_LOG" >&2 + return 1 + fi + if grep -qE "CUDA out of memory|OutOfMemory|RuntimeError|Killed" "$SERVER_LOG"; then + echo "Server hit a fatal error:" >&2 + tail -100 "$SERVER_LOG" >&2 + return 1 + fi + if (( SECONDS - START > TIMEOUT )); then + echo "Timeout after ${TIMEOUT}s waiting for server" >&2 + return 1 + fi + sleep 5 + done + echo "Server ready after $((SECONDS - START))s" +} + +stop_server() { + if [[ -n "$SERVER_PID" ]] && kill -0 "$SERVER_PID" 2>/dev/null; then + echo "Stopping ts serve (pgid $SERVER_PID)..." + kill -TERM -"$SERVER_PID" 2>/dev/null || true + for _ in {1..20}; do + kill -0 "$SERVER_PID" 2>/dev/null || break + sleep 1 + done + kill -KILL -"$SERVER_PID" 2>/dev/null || true + fi + SERVER_PID= +} + +wait_for_port_free() { + local port=${1:-8000} + local timeout=${2:-90} + local start=$SECONDS + while ! python3 -c "import socket; s=socket.socket(); s.bind(('127.0.0.1', $port)); s.close()" 2>/dev/null; do + if (( SECONDS - start > timeout )); then + echo "Port ${port} still in use after ${timeout}s" >&2 + return 1 + fi + sleep 1 + done +} + +trap stop_server EXIT # safety net for Ctrl-C / errors + +# Preflight: bail out if port 8000 is already in use +wait_for_port_free 8000 + +SWEEP_TS=$(date +%Y%m%d_%H%M%S) +SWEEP_DIR="${SCRIPT_DIR}/outputs/${SWEEP_TS}" +echo "Sweep outputs: ${SWEEP_DIR}" + +for CONFIG in "${CONFIGS[@]}"; do + echo "=== Running $CONFIG ===" + launch_server "$CONFIG" + + if ! wait_for_ready; then + stop_server + exit 1 + fi + + echo "Warmup..." + evalscope perf \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --url http://127.0.0.1:8000/v1/chat/completions \ + --api openai \ + --dataset swe_smith \ + --dataset-path agentic_dataset.json \ + --max-tokens 500 \ + --multi-turn \ + --number 2 \ + --parallel 2 \ + --extra-args '{"ignore_eos": true}' \ + --dataset-offset 68 \ + --outputs-dir /tmp/outputs + + echo "Benchmark..." + evalscope perf \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --url http://127.0.0.1:8000/v1/chat/completions \ + --api openai \ + --dataset swe_smith \ + --dataset-path agentic_dataset.json \ + --max-tokens 500 \ + --multi-turn \ + --number 4 8 8 16 32 \ + --parallel 1 2 4 8 16 \ + --extra-args '{"ignore_eos": true}' \ + --name $CONFIG \ + --outputs-dir $SWEEP_DIR \ + --no-timestamp + + stop_server + wait_for_port_free 8000 +done diff --git a/test/agentic_benchmark/kimi_k2.5/tokenspeed/collect_outputs.py b/test/agentic_benchmark/kimi_k2.5/tokenspeed/collect_outputs.py new file mode 100755 index 0000000..dc065fc --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/tokenspeed/collect_outputs.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +"""Collect a sweep's per-config benchmark_summary.json files into one CSV.""" + +import argparse +import csv +import json +import re +import sys +from pathlib import Path + +COLUMNS = [ + "config", + "Conc.", + "Latency (tps/user)", + "Throughput (tps/gpu)", + "Approx Cache Hit", + "Decoded Tok/Iter", +] + + +def num_gpus_from_config(config: str) -> int: + m = re.search(r"attn_(?:tp|dp)(\d+)", config) + if not m: + raise ValueError(f"Cannot infer GPU count from config name: {config}") + return int(m.group(1)) + + +def collect(sweep_dir: Path): + rows = [] + for config_dir in sorted(p for p in sweep_dir.iterdir() if p.is_dir()): + config = config_dir.name + n_gpus = num_gpus_from_config(config) + for run_dir in sorted(config_dir.iterdir()): + summary_path = run_dir / "benchmark_summary.json" + if not summary_path.is_file(): + continue + s = json.loads(summary_path.read_text()) + tpot_ms = s["TPOT (ms)"] + decode_tps_user = 1000.0 / tpot_ms if tpot_ms else 0.0 + tps_gpu = s["Total Throughput (tok/s)"] / n_gpus + rows.append( + { + "config": config, + "Conc.": s["Concurrency"], + "Latency (tps/user)": round(decode_tps_user, 2), + "Throughput (tps/gpu)": round(tps_gpu, 2), + "Approx Cache Hit": round(s["KV Cache Hit Rate (%)"], 2), + "Decoded Tok/Iter": round(s["Decoded Tok/Iter"], 4), + } + ) + rows.sort(key=lambda r: (r["config"], r["Conc."])) + return rows + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("sweep_dir", type=Path, help="e.g. outputs/20260505_152734") + ap.add_argument( + "-o", + "--output", + type=Path, + default=None, + help="CSV output path (default: stdout)", + ) + args = ap.parse_args() + + if not args.sweep_dir.is_dir(): + sys.exit(f"Not a directory: {args.sweep_dir}") + + rows = collect(args.sweep_dir) + out = args.output.open("w", newline="") if args.output else sys.stdout + try: + w = csv.DictWriter(out, fieldnames=COLUMNS) + w.writeheader() + w.writerows(rows) + finally: + if args.output: + out.close() + + +if __name__ == "__main__": + main() diff --git a/test/agentic_benchmark/kimi_k2.5/tokenspeed/configs/attn_dp8_moe_ep8.sh b/test/agentic_benchmark/kimi_k2.5/tokenspeed/configs/attn_dp8_moe_ep8.sh new file mode 100755 index 0000000..122074d --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/tokenspeed/configs/attn_dp8_moe_ep8.sh @@ -0,0 +1,30 @@ +#!/usr/bin/bash + +set -euo pipefail + +exec ts serve \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --data-parallel-size 8 \ + --ep-size 8 \ + --max-model-len 80000 \ + --max-num-seqs 16 \ + --max-prefill-tokens 8192 \ + --chunked-prefill-size 8192 \ + --gpu-memory-utilization 0.8 \ + --disable-cuda-graph-padding \ + --trust-remote-code \ + --attention-backend trtllm_mla \ + --moe-backend flashinfer_trtllm \ + --quantization nvfp4 \ + --kv-cache-dtype fp8 \ + --speculative-algorithm EAGLE3 \ + --speculative-draft-model-path lightseekorg/kimi-k2.5-eagle3-mla \ + --speculative-num-steps 3 \ + --speculative-eagle-topk 1 \ + --speculative-num-draft-tokens 4 \ + --speculative-draft-model-quantization unquant \ + --drafter-attention-backend trtllm_mla \ + --enable-cache-report \ + --host 127.0.0.1 \ + --port 8000 \ + --dist-init-addr 127.0.0.1:4000 diff --git a/test/agentic_benchmark/kimi_k2.5/tokenspeed/configs/attn_dp8_moe_tp8.sh b/test/agentic_benchmark/kimi_k2.5/tokenspeed/configs/attn_dp8_moe_tp8.sh new file mode 100755 index 0000000..401a7a6 --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/tokenspeed/configs/attn_dp8_moe_tp8.sh @@ -0,0 +1,30 @@ +#!/usr/bin/bash + +set -euo pipefail + +exec ts serve \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --data-parallel-size 8 \ + --moe-tp-size 8 \ + --max-model-len 80000 \ + --max-num-seqs 16 \ + --max-prefill-tokens 8192 \ + --chunked-prefill-size 8192 \ + --gpu-memory-utilization 0.8 \ + --disable-cuda-graph-padding \ + --trust-remote-code \ + --attention-backend trtllm_mla \ + --moe-backend flashinfer_trtllm \ + --quantization nvfp4 \ + --kv-cache-dtype fp8 \ + --speculative-algorithm EAGLE3 \ + --speculative-draft-model-path lightseekorg/kimi-k2.5-eagle3-mla \ + --speculative-num-steps 3 \ + --speculative-eagle-topk 1 \ + --speculative-num-draft-tokens 4 \ + --speculative-draft-model-quantization unquant \ + --drafter-attention-backend trtllm_mla \ + --enable-cache-report \ + --host 127.0.0.1 \ + --port 8000 \ + --dist-init-addr 127.0.0.1:4000 diff --git a/test/agentic_benchmark/kimi_k2.5/tokenspeed/configs/attn_tp4_moe_ep4.sh b/test/agentic_benchmark/kimi_k2.5/tokenspeed/configs/attn_tp4_moe_ep4.sh new file mode 100755 index 0000000..fa8401e --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/tokenspeed/configs/attn_tp4_moe_ep4.sh @@ -0,0 +1,29 @@ +#!/usr/bin/bash + +set -euo pipefail + +exec ts serve \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --attn-tp-size 4 \ + --ep-size 4 \ + --max-model-len 80000 \ + --max-num-seqs 16 \ + --max-prefill-tokens 8192 \ + --chunked-prefill-size 8192 \ + --gpu-memory-utilization 0.95 \ + --disable-cuda-graph-padding \ + --trust-remote-code \ + --attention-backend tokenspeed_mla \ + --moe-backend flashinfer_trtllm \ + --quantization nvfp4 \ + --kv-cache-dtype fp8 \ + --speculative-algorithm EAGLE3 \ + --speculative-draft-model-path lightseekorg/kimi-k2.5-eagle3-mla \ + --speculative-num-steps 3 \ + --speculative-eagle-topk 1 \ + --speculative-num-draft-tokens 4 \ + --speculative-draft-model-quantization unquant \ + --drafter-attention-backend tokenspeed_mla \ + --enable-cache-report \ + --host 127.0.0.1 \ + --port 8000 diff --git a/test/agentic_benchmark/kimi_k2.5/tokenspeed/configs/attn_tp4_moe_tp4.sh b/test/agentic_benchmark/kimi_k2.5/tokenspeed/configs/attn_tp4_moe_tp4.sh new file mode 100755 index 0000000..8e39177 --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/tokenspeed/configs/attn_tp4_moe_tp4.sh @@ -0,0 +1,29 @@ +#!/usr/bin/bash + +set -euo pipefail + +exec ts serve \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --attn-tp-size 4 \ + --moe-tp-size 4 \ + --max-model-len 80000 \ + --max-num-seqs 16 \ + --max-prefill-tokens 8192 \ + --chunked-prefill-size 8192 \ + --gpu-memory-utilization 0.95 \ + --disable-cuda-graph-padding \ + --trust-remote-code \ + --attention-backend tokenspeed_mla \ + --moe-backend flashinfer_trtllm \ + --quantization nvfp4 \ + --kv-cache-dtype fp8 \ + --speculative-algorithm EAGLE3 \ + --speculative-draft-model-path lightseekorg/kimi-k2.5-eagle3-mla \ + --speculative-num-steps 3 \ + --speculative-eagle-topk 1 \ + --speculative-num-draft-tokens 4 \ + --speculative-draft-model-quantization unquant \ + --drafter-attention-backend tokenspeed_mla \ + --enable-cache-report \ + --host 127.0.0.1 \ + --port 8000 diff --git a/test/agentic_benchmark/kimi_k2.5/tokenspeed/configs/attn_tp8_moe_ep8.sh b/test/agentic_benchmark/kimi_k2.5/tokenspeed/configs/attn_tp8_moe_ep8.sh new file mode 100755 index 0000000..ba76af0 --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/tokenspeed/configs/attn_tp8_moe_ep8.sh @@ -0,0 +1,29 @@ +#!/usr/bin/bash + +set -euo pipefail + +exec ts serve \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --attn-tp-size 8 \ + --ep-size 8 \ + --max-model-len 80000 \ + --max-num-seqs 16 \ + --max-prefill-tokens 8192 \ + --chunked-prefill-size 8192 \ + --gpu-memory-utilization 0.95 \ + --disable-cuda-graph-padding \ + --trust-remote-code \ + --attention-backend tokenspeed_mla \ + --moe-backend flashinfer_trtllm \ + --quantization nvfp4 \ + --kv-cache-dtype fp8 \ + --speculative-algorithm EAGLE3 \ + --speculative-draft-model-path lightseekorg/kimi-k2.5-eagle3-mla \ + --speculative-num-steps 3 \ + --speculative-eagle-topk 1 \ + --speculative-num-draft-tokens 4 \ + --speculative-draft-model-quantization unquant \ + --drafter-attention-backend tokenspeed_mla \ + --enable-cache-report \ + --host 127.0.0.1 \ + --port 8000 diff --git a/test/agentic_benchmark/kimi_k2.5/tokenspeed/configs/attn_tp8_moe_tp8.sh b/test/agentic_benchmark/kimi_k2.5/tokenspeed/configs/attn_tp8_moe_tp8.sh new file mode 100755 index 0000000..10bc6f1 --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/tokenspeed/configs/attn_tp8_moe_tp8.sh @@ -0,0 +1,29 @@ +#!/usr/bin/bash + +set -euo pipefail + +exec ts serve \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --attn-tp-size 8 \ + --moe-tp-size 8 \ + --max-model-len 80000 \ + --max-num-seqs 16 \ + --max-prefill-tokens 8192 \ + --chunked-prefill-size 8192 \ + --gpu-memory-utilization 0.95 \ + --disable-cuda-graph-padding \ + --trust-remote-code \ + --attention-backend tokenspeed_mla \ + --moe-backend flashinfer_trtllm \ + --quantization nvfp4 \ + --kv-cache-dtype fp8 \ + --speculative-algorithm EAGLE3 \ + --speculative-draft-model-path lightseekorg/kimi-k2.5-eagle3-mla \ + --speculative-num-steps 3 \ + --speculative-eagle-topk 1 \ + --speculative-num-draft-tokens 4 \ + --speculative-draft-model-quantization unquant \ + --drafter-attention-backend tokenspeed_mla \ + --enable-cache-report \ + --host 127.0.0.1 \ + --port 8000 diff --git a/test/agentic_benchmark/kimi_k2.5/trtllm/README.md b/test/agentic_benchmark/kimi_k2.5/trtllm/README.md new file mode 100644 index 0000000..4e5541a --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/trtllm/README.md @@ -0,0 +1,57 @@ +# Agentic Benchmark — TRT-LLM + +Sweep `trtllm-serve` against an agentic, multi-turn workload (SWE-Smith) at a fixed +set of attention/MoE parallelism layouts and report per-config throughput, latency, +and KV-cache hit rate. + +Tested with TRT-LLM `1.3.0rc13`. Server listens on port **8001**. + +## Layout + +``` +agentic_bench.sh # main sweep: dataset prep -> for config in CONFIGS: launch, wait, bench, kill +configs/ # one yaml per parallelism layout (passed to trtllm-serve via --config) +collect_outputs.py # parse a sweep into a flat CSV +outputs///parallel_

_number_/ # per-run evalscope artifacts +``` + +## Run a sweep + +```bash +cd test/agentic_benchmark/trtllm +./agentic_bench.sh +``` + +The script (1) installs evalscope at the pinned commit, (2) builds the SWE-Smith +multi-turn dataset, (3) iterates each config in `CONFIGS=()`: launch server, poll +`/health` until ready, run `evalscope perf`, kill server, wait for the port to be +free, repeat. Aborts the whole sweep on the first failure (`set -e`). + +To narrow the matrix, comment out entries in the `CONFIGS=()` array. + +## Configs + +Each `configs/*.yaml` is a `trtllm-serve --config` file specifying: + +- `tensor_parallel_size`, `enable_attention_dp`, `moe_expert_parallel_size` +- `max_batch_size`, `cuda_graph_config.batch_sizes` +- `kv_cache_config.dtype` +- `moe_config.backend` (TRTLLM for TP-attn, CUTEDSL for DP-attn) +- `attention_dp_config` (DP variants only) — `enable_balance`, + `enable_kv_cache_aware_routing`, batching/timeout iters +- `speculative_config` — Eagle3 with `lightseekorg/kimi-k2.5-eagle3-mla` + +Naming: `attn__moe_` where `X ∈ {tp4,tp8,dp8}` and `Y ∈ {tp4,tp8,ep4,ep8}`. +World size = the number after `attn_(tp|dp)`. + +## Collect results + +```bash +python3 collect_outputs.py outputs/ -o sweep.csv +``` + +Emits one row per (config, concurrency) with `Conc.`, `Latency (tps/user)`, +`Throughput (tps/gpu)`, `Approx Cache Hit`, `Decoded Tok/Iter`. `tps/gpu` divides +the system-wide `Total Throughput (tok/s)` by the GPU count inferred from the +config name; the other metrics come straight from `benchmark_summary.json` +(same numbers as evalscope's `performance_summary.txt` Request Metrics table). diff --git a/test/agentic_benchmark/kimi_k2.5/trtllm/agentic_bench.sh b/test/agentic_benchmark/kimi_k2.5/trtllm/agentic_bench.sh new file mode 100755 index 0000000..2e18cf3 --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/trtllm/agentic_bench.sh @@ -0,0 +1,154 @@ +#!/usr/bin/bash + +# Tested on TensorRT-LLM tag: v1.3.0rc13 + +set -euo pipefail + +# Prepare dataset +EVALSCOPE_COMMIT=acd09b44384d53174768bb1063f675420f76fae9 +pip install "evalscope[perf] @ git+https://github.com/modelscope/evalscope.git@${EVALSCOPE_COMMIT}" + +[ -f build_swe_smith_dataset.py ] || wget https://raw.githubusercontent.com/modelscope/evalscope/${EVALSCOPE_COMMIT}/examples/perf/build_swe_smith_dataset.py \ + -O build_swe_smith_dataset.py + +# Note: Only 71 conversations can be built +[ -f agentic_dataset.json ] || python3 build_swe_smith_dataset.py \ + --model moonshotai/Kimi-K2.5 \ + --first-turn-length 50000 \ + --subsequent-turn-length 800 \ + --min-turns 10 \ + --max-turns 15 \ + --number 128 \ + --output-path agentic_dataset.json \ + --num-workers 32 + +# Sweep configs +CONFIGS=( + attn_tp4_moe_tp4 + attn_tp4_moe_ep4 + attn_tp8_moe_tp8 + attn_tp8_moe_ep8 + attn_dp8_moe_tp8 + attn_dp8_moe_ep8 +) + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SERVER_PID= +SERVER_LOG= + +launch_server() { + local config=$1 + SERVER_LOG=/tmp/trtllm_server_${config}.log + setsid trtllm-serve nvidia/Kimi-K2.5-NVFP4 \ + --max_num_tokens 8192 \ + --max_seq_len 80000 \ + --enable_chunked_prefill \ + --trust_remote_code \ + --config "${SCRIPT_DIR}/configs/${config}.yaml" \ + --host 127.0.0.1 \ + --port 8001 \ + > "$SERVER_LOG" 2>&1 & + SERVER_PID=$! +} + +wait_for_ready() { + local TIMEOUT=600 + local START=$SECONDS + until curl -sf -o /dev/null http://127.0.0.1:8001/health; do + if ! kill -0 "$SERVER_PID" 2>/dev/null; then + echo "Server died early. Last log lines:" >&2 + tail -100 "$SERVER_LOG" >&2 + return 1 + fi + if grep -qE "CUDA out of memory|OutOfMemory|RuntimeError|Killed" "$SERVER_LOG"; then + echo "Server hit a fatal error:" >&2 + tail -100 "$SERVER_LOG" >&2 + return 1 + fi + if (( SECONDS - START > TIMEOUT )); then + echo "Timeout after ${TIMEOUT}s waiting for server" >&2 + return 1 + fi + sleep 5 + done + echo "Server ready after $((SECONDS - START))s" +} + +stop_server() { + if [[ -n "$SERVER_PID" ]] && kill -0 "$SERVER_PID" 2>/dev/null; then + echo "Stopping trtllm-serve (pgid $SERVER_PID)..." + kill -TERM -"$SERVER_PID" 2>/dev/null || true + for _ in {1..20}; do + kill -0 "$SERVER_PID" 2>/dev/null || break + sleep 1 + done + kill -KILL -"$SERVER_PID" 2>/dev/null || true + fi + SERVER_PID= +} + +wait_for_port_free() { + local port=${1:-8001} + local timeout=${2:-90} + local start=$SECONDS + while ! python3 -c "import socket; s=socket.socket(); s.bind(('127.0.0.1', $port)); s.close()" 2>/dev/null; do + if (( SECONDS - start > timeout )); then + echo "Port ${port} still in use after ${timeout}s" >&2 + return 1 + fi + sleep 1 + done +} + +trap stop_server EXIT # safety net for Ctrl-C / errors + +# Preflight: bail out if port 8001 is already in use +wait_for_port_free 8001 + +SWEEP_TS=$(date +%Y%m%d_%H%M%S) +SWEEP_DIR="${SCRIPT_DIR}/outputs/${SWEEP_TS}" +echo "Sweep outputs: ${SWEEP_DIR}" + +for CONFIG in "${CONFIGS[@]}"; do + echo "=== Running $CONFIG ===" + launch_server "$CONFIG" + + if ! wait_for_ready; then + stop_server + exit 1 + fi + + echo "Warmup..." + evalscope perf \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --url http://127.0.0.1:8001/v1/chat/completions \ + --api openai \ + --dataset swe_smith \ + --dataset-path agentic_dataset.json \ + --max-tokens 500 \ + --multi-turn \ + --number 2 \ + --parallel 2 \ + --extra-args '{"ignore_eos": true}' \ + --dataset-offset 68 \ + --outputs-dir /tmp/outputs + + echo "Benchmark..." + evalscope perf \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --url http://127.0.0.1:8001/v1/chat/completions \ + --api openai \ + --dataset swe_smith \ + --dataset-path agentic_dataset.json \ + --max-tokens 500 \ + --multi-turn \ + --number 4 8 8 16 32 \ + --parallel 1 2 4 8 16 \ + --extra-args '{"ignore_eos": true}' \ + --name $CONFIG \ + --outputs-dir $SWEEP_DIR \ + --no-timestamp + + stop_server + wait_for_port_free 8001 +done diff --git a/test/agentic_benchmark/kimi_k2.5/trtllm/collect_outputs.py b/test/agentic_benchmark/kimi_k2.5/trtllm/collect_outputs.py new file mode 100755 index 0000000..dc065fc --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/trtllm/collect_outputs.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +"""Collect a sweep's per-config benchmark_summary.json files into one CSV.""" + +import argparse +import csv +import json +import re +import sys +from pathlib import Path + +COLUMNS = [ + "config", + "Conc.", + "Latency (tps/user)", + "Throughput (tps/gpu)", + "Approx Cache Hit", + "Decoded Tok/Iter", +] + + +def num_gpus_from_config(config: str) -> int: + m = re.search(r"attn_(?:tp|dp)(\d+)", config) + if not m: + raise ValueError(f"Cannot infer GPU count from config name: {config}") + return int(m.group(1)) + + +def collect(sweep_dir: Path): + rows = [] + for config_dir in sorted(p for p in sweep_dir.iterdir() if p.is_dir()): + config = config_dir.name + n_gpus = num_gpus_from_config(config) + for run_dir in sorted(config_dir.iterdir()): + summary_path = run_dir / "benchmark_summary.json" + if not summary_path.is_file(): + continue + s = json.loads(summary_path.read_text()) + tpot_ms = s["TPOT (ms)"] + decode_tps_user = 1000.0 / tpot_ms if tpot_ms else 0.0 + tps_gpu = s["Total Throughput (tok/s)"] / n_gpus + rows.append( + { + "config": config, + "Conc.": s["Concurrency"], + "Latency (tps/user)": round(decode_tps_user, 2), + "Throughput (tps/gpu)": round(tps_gpu, 2), + "Approx Cache Hit": round(s["KV Cache Hit Rate (%)"], 2), + "Decoded Tok/Iter": round(s["Decoded Tok/Iter"], 4), + } + ) + rows.sort(key=lambda r: (r["config"], r["Conc."])) + return rows + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("sweep_dir", type=Path, help="e.g. outputs/20260505_152734") + ap.add_argument( + "-o", + "--output", + type=Path, + default=None, + help="CSV output path (default: stdout)", + ) + args = ap.parse_args() + + if not args.sweep_dir.is_dir(): + sys.exit(f"Not a directory: {args.sweep_dir}") + + rows = collect(args.sweep_dir) + out = args.output.open("w", newline="") if args.output else sys.stdout + try: + w = csv.DictWriter(out, fieldnames=COLUMNS) + w.writeheader() + w.writerows(rows) + finally: + if args.output: + out.close() + + +if __name__ == "__main__": + main() diff --git a/test/agentic_benchmark/kimi_k2.5/trtllm/configs/attn_dp8_moe_ep8.yaml b/test/agentic_benchmark/kimi_k2.5/trtllm/configs/attn_dp8_moe_ep8.yaml new file mode 100644 index 0000000..8c9b32b --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/trtllm/configs/attn_dp8_moe_ep8.yaml @@ -0,0 +1,21 @@ +tensor_parallel_size: 8 +enable_attention_dp: true +moe_expert_parallel_size: 8 +max_batch_size: 16 +cuda_graph_config: + batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16] +kv_cache_config: + dtype: fp8 + free_gpu_memory_fraction: 0.8 +attention_dp_config: + enable_balance: true + batching_wait_iters: 10 + timeout_iters: 500 + enable_kv_cache_aware_routing: true +moe_config: + backend: CUTEDSL +speculative_config: + decoding_type: Eagle3 + max_draft_len: 3 + speculative_model_dir: lightseekorg/kimi-k2.5-eagle3-mla + eagle3_one_model: true diff --git a/test/agentic_benchmark/kimi_k2.5/trtllm/configs/attn_dp8_moe_tp8.yaml b/test/agentic_benchmark/kimi_k2.5/trtllm/configs/attn_dp8_moe_tp8.yaml new file mode 100644 index 0000000..94dc1d0 --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/trtllm/configs/attn_dp8_moe_tp8.yaml @@ -0,0 +1,21 @@ +tensor_parallel_size: 8 +enable_attention_dp: true +moe_expert_parallel_size: 1 +max_batch_size: 16 +cuda_graph_config: + batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16] +kv_cache_config: + dtype: fp8 + free_gpu_memory_fraction: 0.8 +attention_dp_config: + enable_balance: true + batching_wait_iters: 10 + timeout_iters: 500 + enable_kv_cache_aware_routing: true +moe_config: + backend: CUTEDSL +speculative_config: + decoding_type: Eagle3 + max_draft_len: 3 + speculative_model_dir: lightseekorg/kimi-k2.5-eagle3-mla + eagle3_one_model: true diff --git a/test/agentic_benchmark/kimi_k2.5/trtllm/configs/attn_tp4_moe_ep4.yaml b/test/agentic_benchmark/kimi_k2.5/trtllm/configs/attn_tp4_moe_ep4.yaml new file mode 100644 index 0000000..2de4da1 --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/trtllm/configs/attn_tp4_moe_ep4.yaml @@ -0,0 +1,15 @@ +tensor_parallel_size: 4 +enable_attention_dp: false +moe_expert_parallel_size: 4 +max_batch_size: 16 +cuda_graph_config: + batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16] +kv_cache_config: + dtype: fp8 +moe_config: + backend: TRTLLM +speculative_config: + decoding_type: Eagle3 + max_draft_len: 3 + speculative_model_dir: lightseekorg/kimi-k2.5-eagle3-mla + eagle3_one_model: true diff --git a/test/agentic_benchmark/kimi_k2.5/trtllm/configs/attn_tp4_moe_tp4.yaml b/test/agentic_benchmark/kimi_k2.5/trtllm/configs/attn_tp4_moe_tp4.yaml new file mode 100644 index 0000000..c5d08a9 --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/trtllm/configs/attn_tp4_moe_tp4.yaml @@ -0,0 +1,15 @@ +tensor_parallel_size: 4 +enable_attention_dp: false +moe_expert_parallel_size: 1 +max_batch_size: 16 +cuda_graph_config: + batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16] +kv_cache_config: + dtype: fp8 +moe_config: + backend: TRTLLM +speculative_config: + decoding_type: Eagle3 + max_draft_len: 3 + speculative_model_dir: lightseekorg/kimi-k2.5-eagle3-mla + eagle3_one_model: true diff --git a/test/agentic_benchmark/kimi_k2.5/trtllm/configs/attn_tp8_moe_ep8.yaml b/test/agentic_benchmark/kimi_k2.5/trtllm/configs/attn_tp8_moe_ep8.yaml new file mode 100644 index 0000000..8334251 --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/trtllm/configs/attn_tp8_moe_ep8.yaml @@ -0,0 +1,15 @@ +tensor_parallel_size: 8 +enable_attention_dp: false +moe_expert_parallel_size: 8 +max_batch_size: 16 +cuda_graph_config: + batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16] +kv_cache_config: + dtype: fp8 +moe_config: + backend: TRTLLM +speculative_config: + decoding_type: Eagle3 + max_draft_len: 3 + speculative_model_dir: lightseekorg/kimi-k2.5-eagle3-mla + eagle3_one_model: true diff --git a/test/agentic_benchmark/kimi_k2.5/trtllm/configs/attn_tp8_moe_tp8.yaml b/test/agentic_benchmark/kimi_k2.5/trtllm/configs/attn_tp8_moe_tp8.yaml new file mode 100644 index 0000000..24de709 --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/trtllm/configs/attn_tp8_moe_tp8.yaml @@ -0,0 +1,15 @@ +tensor_parallel_size: 8 +enable_attention_dp: false +moe_expert_parallel_size: 1 +max_batch_size: 16 +cuda_graph_config: + batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16] +kv_cache_config: + dtype: fp8 +moe_config: + backend: TRTLLM +speculative_config: + decoding_type: Eagle3 + max_draft_len: 3 + speculative_model_dir: lightseekorg/kimi-k2.5-eagle3-mla + eagle3_one_model: true diff --git a/test/agentic_benchmark/kimi_k2.5/vllm/agentic_bench.sh b/test/agentic_benchmark/kimi_k2.5/vllm/agentic_bench.sh new file mode 100755 index 0000000..107cbf2 --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/vllm/agentic_bench.sh @@ -0,0 +1,142 @@ +#!/usr/bin/bash + +# Tested on vLLM tag: 0.22.0 + +set -euo pipefail + +# Prepare dataset +EVALSCOPE_COMMIT=acd09b44384d53174768bb1063f675420f76fae9 +pip install "evalscope[perf] @ git+https://github.com/modelscope/evalscope.git@${EVALSCOPE_COMMIT}" + +[ -f build_swe_smith_dataset.py ] || wget https://raw.githubusercontent.com/modelscope/evalscope/${EVALSCOPE_COMMIT}/examples/perf/build_swe_smith_dataset.py \ + -O build_swe_smith_dataset.py + +# Note: Only 71 conversations can be built +[ -f agentic_dataset.json ] || python3 build_swe_smith_dataset.py \ + --model moonshotai/Kimi-K2.5 \ + --first-turn-length 50000 \ + --subsequent-turn-length 800 \ + --min-turns 10 \ + --max-turns 15 \ + --number 128 \ + --output-path agentic_dataset.json \ + --num-workers 32 + +# Sweep configs +CONFIGS=( + attn_tp4_moe_tp4 + attn_tp4_moe_ep4 +) + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SERVER_PID= +SERVER_LOG= + +launch_server() { + local config=$1 + SERVER_LOG=/tmp/vllm_server_${config}.log + setsid ${SCRIPT_DIR}/configs/${config}.sh > "$SERVER_LOG" 2>&1 & + SERVER_PID=$! +} + +wait_for_ready() { + local TIMEOUT=1200 + local START=$SECONDS + until curl -sf -o /dev/null http://127.0.0.1:8002/health; do + if ! kill -0 "$SERVER_PID" 2>/dev/null; then + echo "Server died early. Last log lines:" >&2 + tail -100 "$SERVER_LOG" >&2 + return 1 + fi + if grep -qE "CUDA out of memory|OutOfMemory|RuntimeError|Killed" "$SERVER_LOG"; then + echo "Server hit a fatal error:" >&2 + tail -100 "$SERVER_LOG" >&2 + return 1 + fi + if (( SECONDS - START > TIMEOUT )); then + echo "Timeout after ${TIMEOUT}s waiting for server" >&2 + return 1 + fi + sleep 5 + done + echo "Server ready after $((SECONDS - START))s" +} + +stop_server() { + if [[ -n "$SERVER_PID" ]] && kill -0 "$SERVER_PID" 2>/dev/null; then + echo "Stopping vllm serve (pgid $SERVER_PID)..." + kill -TERM -"$SERVER_PID" 2>/dev/null || true + for _ in {1..20}; do + kill -0 "$SERVER_PID" 2>/dev/null || break + sleep 1 + done + kill -KILL -"$SERVER_PID" 2>/dev/null || true + fi + SERVER_PID= +} + +wait_for_port_free() { + local port=${1:-8002} + local timeout=${2:-90} + local start=$SECONDS + while ! python3 -c "import socket; s=socket.socket(); s.bind(('127.0.0.1', $port)); s.close()" 2>/dev/null; do + if (( SECONDS - start > timeout )); then + echo "Port ${port} still in use after ${timeout}s" >&2 + return 1 + fi + sleep 1 + done +} + +trap stop_server EXIT # safety net for Ctrl-C / errors + +# Preflight: bail out if port 8002 is already in use +wait_for_port_free 8002 + +SWEEP_TS=$(date +%Y%m%d_%H%M%S) +SWEEP_DIR="${SCRIPT_DIR}/outputs/${SWEEP_TS}" +echo "Sweep outputs: ${SWEEP_DIR}" + +for CONFIG in "${CONFIGS[@]}"; do + echo "=== Running $CONFIG ===" + launch_server "$CONFIG" + + if ! wait_for_ready; then + stop_server + exit 1 + fi + + echo "Warmup..." + evalscope perf \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --url http://127.0.0.1:8002/v1/chat/completions \ + --api openai \ + --dataset swe_smith \ + --dataset-path agentic_dataset.json \ + --max-tokens 500 \ + --multi-turn \ + --number 2 \ + --parallel 2 \ + --extra-args '{"ignore_eos": true}' \ + --dataset-offset 68 \ + --outputs-dir /tmp/outputs + + echo "Benchmark..." + evalscope perf \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --url http://127.0.0.1:8002/v1/chat/completions \ + --api openai \ + --dataset swe_smith \ + --dataset-path agentic_dataset.json \ + --max-tokens 500 \ + --multi-turn \ + --number 4 8 8 16 32 \ + --parallel 1 2 4 8 16 \ + --extra-args '{"ignore_eos": true}' \ + --name $CONFIG \ + --outputs-dir $SWEEP_DIR \ + --no-timestamp + + stop_server + wait_for_port_free 8002 +done diff --git a/test/agentic_benchmark/kimi_k2.5/vllm/collect_outputs.py b/test/agentic_benchmark/kimi_k2.5/vllm/collect_outputs.py new file mode 100755 index 0000000..dc065fc --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/vllm/collect_outputs.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python3 +"""Collect a sweep's per-config benchmark_summary.json files into one CSV.""" + +import argparse +import csv +import json +import re +import sys +from pathlib import Path + +COLUMNS = [ + "config", + "Conc.", + "Latency (tps/user)", + "Throughput (tps/gpu)", + "Approx Cache Hit", + "Decoded Tok/Iter", +] + + +def num_gpus_from_config(config: str) -> int: + m = re.search(r"attn_(?:tp|dp)(\d+)", config) + if not m: + raise ValueError(f"Cannot infer GPU count from config name: {config}") + return int(m.group(1)) + + +def collect(sweep_dir: Path): + rows = [] + for config_dir in sorted(p for p in sweep_dir.iterdir() if p.is_dir()): + config = config_dir.name + n_gpus = num_gpus_from_config(config) + for run_dir in sorted(config_dir.iterdir()): + summary_path = run_dir / "benchmark_summary.json" + if not summary_path.is_file(): + continue + s = json.loads(summary_path.read_text()) + tpot_ms = s["TPOT (ms)"] + decode_tps_user = 1000.0 / tpot_ms if tpot_ms else 0.0 + tps_gpu = s["Total Throughput (tok/s)"] / n_gpus + rows.append( + { + "config": config, + "Conc.": s["Concurrency"], + "Latency (tps/user)": round(decode_tps_user, 2), + "Throughput (tps/gpu)": round(tps_gpu, 2), + "Approx Cache Hit": round(s["KV Cache Hit Rate (%)"], 2), + "Decoded Tok/Iter": round(s["Decoded Tok/Iter"], 4), + } + ) + rows.sort(key=lambda r: (r["config"], r["Conc."])) + return rows + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("sweep_dir", type=Path, help="e.g. outputs/20260505_152734") + ap.add_argument( + "-o", + "--output", + type=Path, + default=None, + help="CSV output path (default: stdout)", + ) + args = ap.parse_args() + + if not args.sweep_dir.is_dir(): + sys.exit(f"Not a directory: {args.sweep_dir}") + + rows = collect(args.sweep_dir) + out = args.output.open("w", newline="") if args.output else sys.stdout + try: + w = csv.DictWriter(out, fieldnames=COLUMNS) + w.writeheader() + w.writerows(rows) + finally: + if args.output: + out.close() + + +if __name__ == "__main__": + main() diff --git a/test/agentic_benchmark/kimi_k2.5/vllm/configs/attn_tp4_moe_ep4.sh b/test/agentic_benchmark/kimi_k2.5/vllm/configs/attn_tp4_moe_ep4.sh new file mode 100755 index 0000000..583aec5 --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/vllm/configs/attn_tp4_moe_ep4.sh @@ -0,0 +1,20 @@ +#!/usr/bin/bash + +set -euo pipefail + +exec vllm serve \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --tensor-parallel-size 4 \ + --enable-expert-parallel \ + --max-model-len 80000 \ + --max-num-seqs 16 \ + --max-num-batched-tokens 8192 \ + --gpu-memory-utilization 0.9 \ + --trust-remote-code \ + --quantization modelopt_fp4 \ + --kv-cache-dtype fp8 \ + --attention-backend auto \ + --moe-backend flashinfer_trtllm \ + --speculative-config '{"method":"eagle3","model":"lightseekorg/kimi-k2.5-eagle3-mla","num_speculative_tokens":3}' \ + --host 127.0.0.1 \ + --port 8002 diff --git a/test/agentic_benchmark/kimi_k2.5/vllm/configs/attn_tp4_moe_tp4.sh b/test/agentic_benchmark/kimi_k2.5/vllm/configs/attn_tp4_moe_tp4.sh new file mode 100755 index 0000000..0058206 --- /dev/null +++ b/test/agentic_benchmark/kimi_k2.5/vllm/configs/attn_tp4_moe_tp4.sh @@ -0,0 +1,19 @@ +#!/usr/bin/bash + +set -euo pipefail + +exec vllm serve \ + --model nvidia/Kimi-K2.5-NVFP4 \ + --tensor-parallel-size 4 \ + --max-model-len 80000 \ + --max-num-seqs 16 \ + --max-num-batched-tokens 8192 \ + --gpu-memory-utilization 0.9 \ + --trust-remote-code \ + --quantization modelopt_fp4 \ + --kv-cache-dtype fp8 \ + --attention-backend auto \ + --moe-backend flashinfer_trtllm \ + --speculative-config '{"method":"eagle3","model":"lightseekorg/kimi-k2.5-eagle3-mla","num_speculative_tokens":3}' \ + --host 127.0.0.1 \ + --port 8002 diff --git a/test/ci/README.md b/test/ci/README.md new file mode 100644 index 0000000..69d6967 --- /dev/null +++ b/test/ci/README.md @@ -0,0 +1,75 @@ +# CI Task Specs + +`test/ci/` is the source of truth for CI task declarations consumed by +`test/ci_system/pipeline.py`. + +Current trigger values: + +- `per-commit` +- `manual` +- `nightly` + +Supported task types: + +- `ut` +- `server_smoke` +- `eval` +- `perf` + +Currently configured task directories: + +- `eval` +- `ut` + +Each task expands into one matrix entry per runner label. Add a top-level +`priority` to a task YAML to bias dispatch order. GitHub Actions starts matrix +jobs in include-list order, so `high` entries reach a contended runner pool +before `normal` (the default) and `low`. Tasks that omit `priority` keep their +original ordering. + +`priority` accepts either a scalar (applies to every label of the task) or a +per-label mapping (only the listed labels are overridden; every other label +stays at `normal`): + +```yaml +# whole task at high +priority: high + +# only the b300-1gpu instance drops to low; h100-1gpu / b200-1gpu / ... +# of the same task keep the default normal +priority: + b300-1gpu: low +``` + +Typical use: lower a 1gpu kernel unit-test on `b300-1gpu` so the heavier +b300-4gpu evals that share the same box claim the runner first, without +disturbing the same task's ordering on the other GPU families. + +`optional` marks a task or per-label matrix entry as non-blocking. +Optional entries are emitted with `matrix.optional: true`, and the PR workflows +map that to GitHub Actions `continue-on-error`. + +```yaml +# whole task can fail without blocking the workflow +optional: true + +# only the MI355 bench entry is non-blocking; the MI350 entry of the same +# task still blocks on failure +optional: + amd-mi355-1gpu-bench: true +``` + +`b200-` labels are the default B200 runners. Set the +`TOKENSPEED_B200_RUNNER_LABEL` repository variable in GitHub Actions +(`Settings` -> `Secrets and variables` -> `Actions` -> `Variables`) to a +non-empty runner family such as `b200v2` to temporarily route them to +`b200v2-` without editing task YAML. Leave the variable unset or empty to +use the default `b200-` labels. +The CI system derives `SM` from common runner label prefixes by default: +`h100`/`h200` use `sm90`, `b200`/`gb200` use `sm100`, and `b300`/`gb300` use +`sm103`. Use `runner.env.